Rockwell P1036 Select the number DFS

P1036 selection

code

/*Siberian Squirrel*/
/*Cute JinFish*/
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const double PI = acos(-1), eps = 1e-8;
/*const int MOD = 998244353, r = 119, k = 23, g = 3;
const int MOD = 1004535809, r = 479, k = 21, g = 3;*/
const int INF = 0x3f3f3f3f, MOD = 1e9 + 7;
const int M = 1e7 + 10, N = 5e6 + 10;
int sgn(double x) {
    
    
    if(fabs(x) < eps) return 0;
    return x < 0? -1: 1;
}
//inline int rnd(){static int seed=2333;return seed=(((seed*666666ll+20050818)%998244353)^1000000007)%1004535809;}
//double Rand() {return (double)rand() / RAND_MAX;}

int n, k;
int a[N], ans;
int prime[N], cnt;
bool vis[N];

ll quick_pow(ll ans, ll p, ll res = 1) {
    
    
    for(; p; p >>= 1, ans = ans * ans % MOD)
        if(p & 1) res = res * ans % MOD;
    return res % MOD;
}

void init() {
    
    
    for(int i = 2; i < N; ++ i) {
    
    
        if(!vis[i]) prime[++ cnt] = i;
        for(int j = 1; j <= cnt && prime[j] * i < N; ++ j) {
    
    
            vis[i * prime[j]] = true;
            if(i % prime[j] == 0) break;
        }
    }
}

void dfs(int num, int s, int st) {
    
    
    if(num == k) {
    
    
        if(!vis[s]) ans ++;
        return;
    }
    for(int i = st; i <= n; ++ i) {
    
    
        dfs(num + 1, s + a[i], i + 1);
    }
    return;
}

void solve(ll res = 0) {
    
    
    dfs(0, 0, 1);
    cout << ans << endl;
}


int main() {
    
    
    ios::sync_with_stdio(false);cin.tie(0);cout.tie(nullptr);
// srand(time(0));
#ifdef ACM_LOCAL
    freopen("input", "r", stdin);
    freopen("output", "w", stdout);
#endif
    init();
    int o = 1;
//	cin >> o;
    while(o --) {
    
    
        cin >> n >> k;
        for(int i = 1; i <= n; ++ i) cin >> a[i];
        solve();
    }
    return 0;
}