yukicoder : No.3257 +|+

ちょっと手間取ったけどまぁ解けた。
https://yukicoder.me/problems/no/3257

問題

1-indexedなN要素の整数列Aが与えられる。
1≦i＜j≦Nとなる(i,j)の対のうち、A[i]+A[j]が(i+j)で割り切れるものは何通りか。

解法

max(A)がO(N)とすると、ある定数Bに対し、i≧Bの場合A[i]+A[j]は(i+j)の高々O(N/B)倍にしかならない。
そこで、B=√N程度として

• i≦Bのケースは愚直に総当たりする
• i＞Bのケースは、A[i]+A[j]=k(i+j)となるkを、1～B程度まで総当たりする。
  • A'[i]= A[i]-k*iとするとA'[i]+A'[j]=0となる(i,j)を求める問題に帰着できるので、k1つあたりO(N)で求められる。

int N;
int A[202020];
unordered_map<int,int> M[401];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i+1];
	}
	
	ll ret=0;
	for(i=1;i<=min(N,2000);i++) {
		for(j=i+1;j<=N;j++) if((A[i]+A[j])%(i+j)==0) ret++;
	}
	
	for(i=2001;i<=N;i++) {
		for(j=1;A[i]-i*j>=-200000;j++) {
			M[j][A[i]-i*j]++;
		}
	}
	for(i=2001;i<=N;i++) {
		for(j=1;A[i]-i*j>=-200000;j++) {
			x=A[i]-i*j;
			M[j][A[i]-i*j]--;
			if(M[j].count(-x)) ret+=M[j][-x];
		}
	}
	cout<<ret<<endl;
	
	
	
	
}

まとめ

ちょっと定数倍が厳しかったかも。