ABC備忘録【AtCoder Beginner Contest 342 D - Square Pair】2

はじめに

この記事を投稿したのですが、もっと簡単に書けるやり方があったのでそちらも残しておきます。

考え方

前回の記事では、平方数である条件を「素因数の指数が全て偶数であること」と扱い、素因数分解を用いて指数が奇数となる素因数のペアを作成することで答えを求めました。
今回は、事前に与えられた数字を先に平方数で割れるだけ割っておいて、その結果が等しいもの同士のペアを作成するやり方で答えを求めます。

参考解説

コード

import sys
from collections import defaultdict
import math


# 平方数で割り続ける
def divide_square(n):
    if n == 1:
        return 1
    for i in range(2, int(math.sqrt(n)) + 1):
        square = i * i
        while n % square == 0:
            n //= square
            if n == 1:
                return 0
    return n


N = int(sys.stdin.readline())
l = list(map(int, sys.stdin.readline().split()))

dictionary = defaultdict(int)
for number in l:
    if number != 0:
        n = divide_square(number)
        dictionary[n] += 1

zero_count = l.count(0)
pair_count = zero_count * (N - zero_count) + zero_count * (zero_count - 1) // 2
for n, count in dictionary.items():
    if n == 1:
        # 1と平方数との積は必ず平方数となる
        pair_count += count * dictionary[0]
    pair_count += count * (count - 1) // 2

print(pair_count)

終わりに

平方数で割っていった商が同じもの同志の積は平方数になるという考えは、言われてみれば当たり前ですけどまるで気付けなかったので勉強になりました。