問題
数列中から次の条件を満たす3つ組の数を数える. 比率 r が与えられていて [1, r, r^2] . 数列はソートされていると仮定する.
入力
数列の長さ n 比率 r
数列 arr
出力
可能な3つ組の数
入力例
6 3
1 3 9 9 27 81
出力例
6
3つ組を数える基礎として、まず2つ組を数えます。
数列の値 arr[i] を大きい値から順番にハッシュ mp1 に入れて数えます(mp1[arr[i]]++)。それと同時に arr[i] * r の値が mp1 に入っていたら, その値を mp2 に加えます(mp2[arr[i]] += mp1[arr[i] * r])。こうすると arr[i], arr[i]*r の組が数えられます。
ハッシュ (unordered_map) は定義されていないときの初期値は 0 になります。
3つ組も同じように arr[i] * r の値が mp2 に入っていたら、その値を mp3 に加えます(mp3[arr[i]] += mp2[arr[i] * r])。
数列の値をすべてチェックしたあとに mp3 の値の総和を求めれば、答えになります。ハッシュのループはイテレータを使います。
long countTriplets(vector<long> arr, long r) {
unordered_map<long, long> mp1, mp2, mp3;
int l = arr.size();
for (int i = l - 1; i >= 0; i--) {
mp3[arr[i]] += mp2[arr[i]*r];
mp2[arr[i]] += mp1[arr[i]*r];
mp1[arr[i]]++;
}
long res = 0;
for (unordered_map<long, long>::iterator itr = mp3.begin(); itr != mp3.end(); itr++) {
res += itr->second;
}
return res;
}