けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 342 D - Square Pair (1Q, 緑色, 400 点)

AtCoder AtCoder400点 ABC-D NoviSteps1Q 緑色diff 数学(整数問題) 条件の言い換え バケット 連想配列(setやmap) 計算量改善:setやmapの活用 集計処理 数列 解空間:O(N^2)個のペア 解空間:O(N^2)通りの選択肢

条件を上手に言い換えよう！

問題概要

N 個の整数 A_{1}, A_{2}, \dots, A_{N} が与えられる。

  • 1 \le i \lt j \le N
  • A_{i} A_{j} が平方数

という条件を満たすような組 (i, j) の個数を求めよ。

制約

  • 2 \le N \le 2 \times 10^{5}
  • 0 \le A_{i} \le 2 \times 10^{5}

考えたこと

これは「条件の言い換え」を頑張る問題。「 A_{i} A_{j} が平方数」という条件を扱いやすいように言い換えよう。これがちょっと難しい。

個人的に思ったのは、平方根を持ち出すとわかりやすかった（ただし、、 A_{i} = 0 の部分は別に考えることとして、 a, b は 1 以上とする）。

  • ab が平方数であるとは、 \sqrt{ab} が整数であることである
  • そのことは、 \sqrt{a}, \sqrt{b} をそれぞれ簡単にしたときに、√ の中身が一致することと同値である

たとえば、 a, b = 18, 50 をかけると ab = 900 となって平方数となるが、これは次のように解釈できる。

  • \sqrt{a} = \sqrt{18} = 3\sqrt{2} である
  • \sqrt{b} = \sqrt{50} = 5\sqrt{2} である
  • これらは \sqrt{2} の部分が一致しているから、かけると整数になる

この言い換えさえできれば、あとは簡単だ。次のように問題は言い換えられる。

\sqrt{A_{1}}, \sqrt{A_{2}}, \dots, \sqrt{A_{N}} をそれぞれ簡単にしたときの √ の中身を B_{1}, B_{2}, \dots, B_{N} とする。

  • 1 \le i \lt j \le N
  • B_{i} = B_{j}

という条件を満たすような組 (i, j) の個数を求めよ。

これは、map などを使って集計処理することで、容易に答えられる。

なお、0 は例外処理する必要がある。0 に何をかけても 0 であって、これは平方数であるからだ。

全体として、 A_{i} の最大値を M として、計算量は O(N (\log N + \sqrt{M}) となる。

コード

#include <bits/stdc++.h>
using namespace std;
const int MAX = 210000;

int main() {
    long long N, A;
    cin >> N;
    vector<long long> nums(MAX, 0);
    for (int i = 0; i < N; i++) {
        cin >> A;
        for (long long j = 2; j * j <= A; j++) {
            while (A % (j * j) == 0) {
                A /= j * j;
            }
        }
        nums[A]++;
    }

    // (0, 他の数) の組の個数
    long long res = N * (N - 1) / 2 - (N - nums[0]) * (N - nums[0] - 1) / 2;

    // 0 を含まない組の個数
    for (long long val = 1; val < MAX; val++) {
        res += nums[val] * (nums[val] - 1) / 2;
    }
    cout << res << endl;
}