yukicoder : No.732 3PrimeCounting

遅刻かつ途中離脱で時間内に解けず。
https://yukicoder.me/problems/no/732

問題

整数Nが与えられる。
N以下の異なる素数a＜b＜cについて、a+b+cが素数となるのは何通りか。

解法

Nの上限は10^5なのでO(N^2)かかるアルゴリズムだと間に合わない。
ただ素数に限定すればだいぶ少ないので、素数を総当たりしよう。
a+b+cも素数なので、素数を2つcおよびa+b+cを総当たりしよう。
先にa+bの候補を列挙しておけばf(N)をN以下の素数の数とするとO(f(N)^2)程度で収まる。

以下の例では、cを大きくするたびにbをcの1つ前の素数としてa+bを数え上げるということをしている。

const int prime_max = 1000000;
int NP,prime[100000],divp[prime_max];
map<int,int> M;

void cprime() {
	if(NP) return;
	for(int i=2;i<prime_max;i++) if(divp[i]==0) {
		//M[i]=NP;
		prime[NP++]=i;
		for(ll j=1LL*i*i;j>=i&&j<prime_max;j+=i) if(divp[j]==0) divp[j]=i;
	}
}

int N;
int cnt[303030];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	cprime();
	
	ll ret=0;
	int a,b,c;
	for(c=2;prime[c]<=N;c++) {
		b=c-1;
		FOR(a,b) cnt[prime[a]+prime[b]]++;
		for(y=c+1;prime[y]<=prime[c]*3;y++) ret+=cnt[prime[y]-prime[c]];
	}
	
	cout<<ret<<endl;
}

まとめ

O(素数の数^2)は想定解だったのか…。