TopCoderOpen 2019 Wildcard Medium Diophantine

これは蟻本読んでてよかったパターン。
https://community.topcoder.com/stat?c=problem_statement&pm=15688

問題

整数N,Dが与えられる。
P[i]を0-originで小さい方からi番目の素数とする。
Q[i]=P[i]*P[i+D]と定義する。

Q[0]～Q[N-1]に含まれる値を重複含め5個選び、その値を昇順にVWXYZとする。
V+W+X+Y=Zを満たす選び方は何通りあるか。

解法

5個の値の加減算を行うので、少なくとも一つは偶数である。
この時、Qの構築法を見ると偶数はP[0]=2を掛けているQ[0]しかない。
Q[0]はQの中で最小値なので、V=Q[0]で決まりである。

後は残り4要素に対し式を満たす要素を抜き出すのだが、これは蟻本序盤にあるくじ引きと同じで、O(N^2*logN)またはunordered_mapでO(N^2)で解ける。
Q[0]+W+Xの組み合わせをmapにしまいつつ、Z-Yを総当たりしよう。

以下のコードはmapだと7秒弱、unordered_mapだと5秒ちょいで通る。

const int prime_max = 1000000;
vector<int> prime;
int NP,divp[prime_max];

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

ll Q[2500];

class Diophantine {
	public:
	long long countSolutions(int N, int D) {
		cprime();
		int i,x,y,z;
		FOR(i,N) Q[i]=1LL*prime[i]*prime[i+D];
		unordered_map<int,int> M;
		ll ret=0;
		FOR(y,N) FOR(x,y+1) M[Q[0]+Q[x]+Q[y]]++;
		for(y=N-1;y>=0;y--) {
			for(z=y+1;z<N;z++) if(M.count(Q[z]-Q[y])) ret+=M[Q[z]-Q[y]];
			for(z=0;z<=y;z++) M[Q[0]+Q[y]+Q[z]]--;
		}
		
		return ret;
	}
	
}

まとめ

Nの上限2000位でもよくない？