《算法笔记》8.1小节——搜索专题-＞深度优先搜索（DFS）问题 C: 【递归入门】组合+判断素数

题目描述

已知 n 个整数b1,b2,…,bn

以及一个整数 k（k＜n）。

从 n 个整数中任选 k 个整数相加，可分别得到一系列的和。

例如当 n=4，k＝3，4 个整数分别为 3，7，12，19 时，可得全部的组合与它们的和为：
　　　　3＋7＋12=22　　3＋7＋19＝29　　7＋12＋19＝38　　3＋12＋19＝34。
　　现在，要求你计算出和为素数共有多少种。

例如上例，只有一种的和为素数：3＋7＋19＝29。

输入

第一行两个整数：n , k （1<=n<=20，k＜n） 
第二行n个整数：x1,x2，…,xn （1<=xi<=5000000） 

输出

一个整数（满足条件的方案数）。 

样例输入

4 3
3 7 12 19

样例输出

1

分析：本质上是B题的强化版本，在B题基础上增加了求和和素数判断。基本上把B的代码拿过来改动一下就可以了。不过下面的代码是递归写的。

#include<algorithm>
#include <iostream>
#include  <cstdlib>
#include  <cstring>
#include   <string>
#include   <vector>
#include   <cstdio>
#include    <queue>
#include    <stack>
#include    <ctime>
#include    <cmath>
#include      <map>
#include      <set>
#define ll long long
#define INF 0x3f3f3f3f
#define db1(x) cout<<#x<<"="<<(x)<<endl
#define db2(x,y) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<endl
#define db3(x,y,z) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<", "<<#z<<"="<<(z)<<endl
#define db4(x,y,z,a) cout<<#x<<"="<<(x)<<", "<<#y<<"="<<(y)<<", "<<#z<<"="<<(z)<<", "<<#a<<"="<<(a)<<endl
using namespace std;

void getprime(int *flag,vector<int>&prime)
{
    flag[0]=flag[1]=1;
    for(int i=2;i<=10000;++i)
        if(!flag[i])
            for(int j=i*i;j<=10000;j=j+i)
                flag[j]=1;
    for(int i=2;i<=10000;++i)
        if(!flag[i])
            prime.push_back(i);

    return;
}

void getans(int n,int k,int index,int *num,int *flag,int *temp,vector<int>prime,vector<int>&ans,int pos)
{
    if(index==k)
    {
        int sum=0,cnt=prime.size();
        for(int i=0;i<k;++i)
            sum+=temp[i];
        int f=0;
        for(int i=0;i<cnt;++i)
        {
            if(sum%prime[i]==0&&sum!=prime[i])
            {
                f=1;break;
            }
        }
        if(!f)ans.push_back(sum);
        return;
    }
    for(int i=pos;i<n;++i)
    {
        if(flag[i]==0)
        {
            temp[index]=num[i],flag[i]=1;
            getans(n,k,index+1,num,flag,temp,prime,ans,i+1);
            flag[i]=0;
        }
    }
    return;
}

int main(void)
{
    #ifdef test
    freopen("in.txt","r",stdin);
    //freopen("out.txt","w",stdout);
    clock_t start=clock();
    #endif //test

    int n,k;scanf("%d%d",&n,&k);
    int num[n+5]={0},flag[20000]={0},temp[k+5]={0};
    vector<int>prime,ans;
    for(int i=0;i<n;++i)
        scanf("%d",&num[i]);
    getprime(flag,prime);
    memset(flag,0,sizeof(int)*(n+5));
    getans(n,k,0,num,flag,temp,prime,ans,0);
    printf("%d\n",ans.size());

    #ifdef test
    clockid_t end=clock();
    double endtime=(double)(end-start)/CLOCKS_PER_SEC;
    printf("\n\n\n\n\n");
    cout<<"Total time:"<<endtime<<"s"<<endl;        //s为单位
    cout<<"Total time:"<<endtime*1000<<"ms"<<endl;    //ms为单位
    #endif //test
    return 0;
}

上面的代码比较繁琐，下面给出AI优化过的代码。

#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>
using namespace std;

#define ll long long

// 判断某个数是否是素数（试除法）
bool is_prime(int num) {
    if (num < 2) return false;
    for (int i = 2; i * i <= num; ++i) {
        if (num % i == 0) return false;
    }
    return true;
}

// 递归搜索所有 k 组合，并判断是否为素数
void dfs(int index, int sum, int count, int n, int k, vector<int>& num, int& result) {
    if (count == k) { // 选满 k 个数
        if (is_prime(sum)) result++;
        return;
    }

    for (int i = index; i < n; ++i) {
        dfs(i + 1, sum + num[i], count + 1, n, k, num, result);
    }
}

int main() {
    int n, k;
    cin >> n >> k;
    vector<int> num(n);
    
    for (int i = 0; i < n; ++i) {
        cin >> num[i];
    }

    int result = 0;
    dfs(0, 0, 0, n, k, num, result);
    cout << result << endl;
    
    return 0;
}