Codeforces Round #Pi Div2 C: Geometric Progression
解法
について考える。
- の時、自分より前に出てきたbの数を記録しておく。
- の時、自分より前に出てきたに記録されている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(); } }