😊

ABC342 D - Square Pairのメモ

2024/02/26に公開

AtCoder

競プロ

tech

コンテストで解けなかったので解説ACです。

考えたこと

まず、組み合わせの問題は片方を固定して考えるのが定石です。全て列挙するのは間に合いませんが、片方を1つ目に固定した時の答えを高速に求める、2つ目に固定した時の答えを... とやっていって最後に足すのは間に合います。今回の問題は数学っぽいので、たぶん計算で求めるのでしょう。

次に、計算で求めるにあたって、そもそも平方数であるとはどういうことなのかを考えました。それについては、素因数分解した時に各素因数が全て偶数個になるはずだと考えました。
（数学は苦手だけど素因数分解はなんか好きなのでそこに思考が行きます…）
逆に、奇数個の素因数を含む場合は平方数ではありません。そして、その奇数個の素因数を掛けたら平方数になるはず。

たとえば12を素因数分解すると $2^2 \times 3^1$ で、3が奇数個なので平方数ではありません。それに3を掛けると $2^2 \times 3^2$ になるので平方数となります。計算すると36で、 $6^2$ なのでたしかに平方数です。
ということは、要素の中に12と3があった場合、12を既に素因数分解していて上記のことを知っていれば、実際に掛け算してみなくとも条件を満たす組み合わせであるとわかります。3の個数を数えておけば、その個数を掛けるだけで高速にカウントするみたいな感じでいけそうな感触を得ました。
まあ、ダメでしたが…

なぜダメだったのか

上記は間違ってはいないのですが、理解が浅かったです。12と3を掛けたら平方数になるのはそうなんですが、別に3でなくとも素因数分解した時に3が余る数なら何でも平方数になります。たとえば12に $2^4 \times 3^1 = 48$ を掛けても平方数になります。 $12 \times 48$ は576ですが、これは $24^2$ なのでたしかに平方数です。
なので、3だろうと12だろうと48だろうと、素因数分解して3が余る数であればどれも同一視できるし、まとめてカウントする必要がありました。
素因数分解して余る数ごとにカウントしたら、その個数から2つ選ぶ組み合わせの数を求めます。個数を $n$ とすると $nC_2$ を計算します。つまり $n \times (n - 1) \div 2$ です。
（ $i<j$ という制約があるので、順序を区別しなくてOK。また、元の並び順は特に結果に影響なし）
これを素因数分解で余った数ごとに求めて足せばOKというわけです。（なお、複数の数が余った場合はかけ合わせます）

ただ、元の値が0の場合は考慮が必要です。0の場合はどの数と組み合わせても条件を満たすので、別途数えて別途足します。
入力例1の[0 3 2 8 12]の場合は、0と各数値との組み合わせがあるので4通りを追加で足します。これは上記の組み合わせの計算とは独立して考えることができます。

実装

fun main() {
    val n = readln().toInt()
    val a = readln().split(" ").map { it.toLong() }

    val map = mutableMapOf<Long, Long>()

    val pfs = a.filter { it != 0L }.map { it to primeFactorial(it) }
    val zeroCount = n - pfs.size

    for(i in pfs.indices) {
        if(pfs[i].first == 1L) {
            map[1] = map.getOrDefault(1, 0) + 1
            continue
        }

        val pf = pfs[i].second

        var m = 1L
        pf.forEach {
            if(it.second % 2 != 0L) {
                m *= it.first
            }
        }

        map[m] = map.getOrDefault(m, 0) + 1
    }

    var ans = 0L
    for(i in map.values) {
        ans += i * (i - 1) / 2
    }
    var tmp = n
    for(i in 1..zeroCount) {
        tmp--
        ans += tmp
    }

    println(ans)
}

fun primeFactorial(paramN: Long): List<Pair<Long, Long>> {
    val result = mutableListOf<Pair<Long, Long>>()
    var n = paramN
    var a = 2L
    while(a * a <= n) {
        if(n % a != 0L) {
            a++
            continue
        }
        var ex = 0L

        while (n % a == 0L) {
            ex++
            n /= a
        }
        result.add(a to ex)
        a++
    }
    if(n != 1L) {
        result.add(n to 1)
    }
    return result
}

この実装としては、1も特別に考慮しています。素因数分解して余った数をカウントするだけだともともと平方数だった場合の考慮が抜けています。平方数は1を掛けたら平方数になる数と考えることができるので、1をキーとする値としてカウントします。

0がある場合ですが、複数ある場合を考えると上記のようなカウント方法になります。ちょっと計算式で考えてもわからなかったので愚直にカウントしました…

計算量としては、素因数分解が $O(\sqrtA_i)$ なので全体で $O(N\sqrtA_i)$ でしょうか。若干心配になったけど大丈夫でした。

考えたこと

なぜダメだったのか

実装

Discussion