CFR Educational 12 D. Simple Subset ( Math + Adhoc )

Problem - D - Codeforces
2 elements with same value should never co-exist since it will be an even number, which will not be a prime number when summed, unless they were both 1. 2 even numbers nor 2 odd numbers can not co-exist as well. So if we make the original input unique, and find valid set with DFS, it will not go any farther than depth of 2, which leads to the time complexity be O ( n ^ 2 ). Preprocess prime table to make it possible to check if some number is prime in constant time.

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1e3 + 3;
const int MAXA = 1e6 + 6;

int not_prime[ MAXA << 1 ];
void prePrime(){
    not_prime[ 1 ] = 1;
    for( int i = 2; i < MAXA * 2; ++i ){
        if( !not_prime[ i ] )
            for( int j = i * 2; j < MAXA * 2; j += i )
                not_prime[ j ] = 1;
    }
}

int n;
int a[ MAXN ];

vector< int > ans;

void dfs( int c, vector< int > &has ){
    if( has.size() > ans.size() )
        ans = has;
    if( c == n ) return;
    bool can_take = true;
    for( int v : has ){
        if( not_prime[ v + a[ c ] ] )
            can_take = false;
    }
    if( can_take )
        has.push_back( a[ c ] ),
        dfs( c + 1, has ),
        has.pop_back();
    dfs( c + 1, has );
}

void solve(){
    int cnt1 = count( a, a + n, 1 );
    sort( a, a + n );
    n = unique( a, a + n ) - a;
    if( a[ 0 ] == 1 ){
        vector< int > vec;
        for( int i = 0; i < cnt1; ++i )
            vec.push_back( 1 );
        dfs( 1, vec );
        vec.clear();
        dfs( 1, vec );
    } else{
        vector< int > vec;
        dfs( 0, vec );
    }
    printf( "%d\n", ans.size() );
    for( int i = 0; i < ans.size(); ++i )
        printf( "%d%c", ans[ i ], i == ans.size() - 1 ? '\n' : ' ' );
}

int main(){
    prePrime();
    scanf( "%d", &n );
    for( int i = 0; i < n; ++i )
        scanf( "%d", a + i );
    solve();
    return 0;
}