Codeforces #875 : Div1 B. The BOSS Can Count Pairs

systestで落として微妙に終わった回。
https://codeforces.com/contest/1830/problem/B

問題

N要素の整数対(A[i],B[i])が与えられる。
2つの整数対(A[i],B[i]),(A[j],B[j])に対しA[i]*A[j]=B[i]+B[j]となる(i,j)の対は何通りか。

解法

A[i]を総当たりする。
A[i]の倍数を総当たりし、B[i]+B[j]となるパターンを数えていこう。

ll T,N,A[202020],B[202020];

unordered_map<int,int> M[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		FOR(i,N+1) M[i].clear();
		FOR(i,N) cin>>A[i];
		FOR(i,N) {
			cin>>B[i];
			M[A[i]][B[i]]++;
		}
		ll ret=0;
		for(i=1;i<=N;i++) if(M[i].size()) {
			for(j=i;j<=N&&1LL*i*j<=2*N;j++) if(M[j].size()) {
				if(i==j) {
					FORR2(a,b,M[i]) if(a<=i*i-a&&M[i].count(i*i-a)) {
						if(a==i*i-a) {
							ret+=1LL*b*(b-1)/2;
						}
						else {
							ret+=1LL*b*M[i][i*i-a];
						}
					}
				}
				else {
					x=i,y=j;
					if(M[i].size()>M[j].size()) swap(x,y);
					FORR2(a,b,M[x]) {
						k=i*j-a;
						if(k>=0&&k<=N&&M[y].count(k)) ret+=1LL*b*M[y][k];
					}
				}
			}
		}
		cout<<ret<<endl;
	}
}

まとめ

うまく枝刈りできるかが重要な問題。