Codeforces #314 Div2 C. Geometric Progression

全問解けたと思ったら２問もミスしていた。
正式には回次は#Piっぽいけど、ここでは#314としています。
http://codeforces.com/contest/567/problem/C

問題

N要素の整数列A[i]とKが与えられる。
Aのうち３要素を抜き出したとき、それが公比Kの等比数列となっている、A[i]*K^2=A[j]*K=A[k] (i＜j＜k)となる(i,j,k)の組の数を求めよ。

解法

色々解法がある。
先にA[j]を総当たりし、j未満のindexにA[j]/Kが登場する数とj以降のindexにA[j]*Kが登場する数を掛け合わせてもよい。
iを順番に見ていってA[i]が初項として登場するケース、２項目として登場するケース、３項目として登場するケースをDPの要領で数え上げてもよい。
以下の回答は後者。

int N,K;
ll A[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) cin>>A[i];
	
	ll ret=0;
	unordered_map<ll,ll> M[2];
	
	for(i=N-1;i>=0;i--) {
		ret +=M[1][A[i]*K];
		M[1][A[i]]+=M[0][A[i]*K];
		M[0][A[i]]++;
	}
	
	cout<<ret<<endl;
}

まとめ

オーバーフローやらかしてしまった。