Codeforces #572 Div1 B. Count Pairs

うーん数学ゲー。
http://codeforces.com/contest/1188/problem/B

問題

N個の整数X[i]と、素数P・整数Kが与えられる。
(X[i]+X[j])*(X[i]^2+X[j]^2) ≡ K (mod P)となるような(i,j)の対は何通りか。

解法

の両辺にを掛けると、となり、となる。
よって、(X[i]^4-X[i]*K)の値が一致するものを数え上げればよい。

int N,P,K;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>P>>K;
	map<ll,int> M;
	ll ret=0;
	FOR(i,N) {
		ll v;
		cin>>v;
		ll w=((v*v%P*v%P*v%P-K*v)%P+P)%P;
		ret+=M[w];
		M[w]++;
	}
	
	cout<<ret<<endl;
	
}

まとめ

(a+b)で括ることにこだわり過ぎた…。