Codeforces Round #Pi Div2 C: Geometric Progression - 宇宙ツイッタラーXの憂鬱

問題

<a href="http://codeforces.com/contest/567/problem/C">Problem - C - Codeforces</a>codeforces.com

解法

$a_i$ について考える。

$a_i=b \times k$ の時、自分より前に出てきたbの数を記録しておく。
$a_i=b \times k^2$ の時、自分より前に出てきた $b \times k$ に記録されているbの数の和を記録しておく。
$a_i=b$ の時、それまでに出てきたbの数に1加える。

コード

import java.util.Arrays;
import java.util.HashMap;
import java.util.Scanner;

public class Main {

	public void solve() {
		Scanner scanner = new Scanner(System.in);
		int N = scanner.nextInt();
		int K = scanner.nextInt();
		int[] a = new int[N];
		for (int i = 0; i < N; i++) {
			a[i] = scanner.nextInt();
		}
		scanner.close();

		if (K == 1) {
			Arrays.sort(a);

			long ans = 0;
			long cnt = 0;
			for (int i = 0; i < N; i++) {
				cnt++;
				if (i == N - 1 || a[i] != a[i + 1]) {
					if (cnt >= 3) {
						long tmp = 1;
						tmp *= cnt * (cnt - 1) * (cnt - 2);
						tmp /= 6;
						ans += tmp;
					}
					cnt = 0;
				}
			}
			System.out.println(ans);
			return;
		}

		long ans = 0;
		long zeroCnt = 0;
		HashMap<Integer, Long> countMap = new HashMap<>();
		HashMap<Integer, Long> dimMap = new HashMap<>();
		for (int i = 0; i < N; i++) {
			// ゼロの時は別途対応
			if (a[i] == 0) {
				zeroCnt++;
				continue;
			}

			// 一番下の時
			if (countMap.containsKey(a[i])) {
				countMap.put(a[i], countMap.get(a[i]) + 1);
			} else {
				countMap.put(a[i], 1L);
			}

			// 真ん中の時
			if (a[i] % K != 0) {
				continue;
			}
			int child = a[i] / K;
			if (dimMap.containsKey(a[i]) && countMap.containsKey(child)) {
				dimMap.put(a[i], dimMap.get(a[i]) + countMap.get(child));
			} else if (countMap.containsKey(child)) {
				dimMap.put(a[i], countMap.get(child));
			}

			// 一番上の時
			if (dimMap.containsKey(child)) {
				ans += dimMap.get(child);
			}
		}

		// ゼロの中から3つゼロを取り出す組み合わせ
		long tmp = zeroCnt * (zeroCnt - 1) * (zeroCnt - 2);
		tmp /= 6;
		ans += tmp;

		System.out.println(ans);

	}

	public static void main(String[] args) {
		new Main().solve();
	}

}