ABC342 D - Square Pair

atcoder.jp

問題

• 長さNの非負整数列Aが与えられるよ
• 以下の条件を満たす整数組$(i,j)$は何個ある？
  • $1 \leq i < j \leq N$
  • $A_i \times A_j$は平方数

成約

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

思考

• iとjを全部見ていくのはまぁ間に合いません…
• オチとしては、iを集計しつつjを1つずつ見ていく、みたいな話
• 平方数の判定
  • $n$を素因数分解して、$a^{i} \times b^{j} \times c^{k} ...$の形にする
  • 指数？i,j,kを2で割ったあまりが、平方数で割ったあまりになる
    • 偶数個ある素因数は、$a \times a$みたいに平方数として取れるのでそれはそう
• 0の取り扱い
  • 0は何をかけても0で、今回の平方数の条件に一致するので、ここのさばき方だけ注意すること
  • 逆に、平方数で割ったあまりがなんであっても、ゼロと掛けると0＝平方数になることも注意

コード

from collections import defaultdict

N = int(input())
A_L = list(map(int, input().split()))

# 素因数分解
def factorization(n):
    arr = []
    tmp = n
    for i in range(2, int(-(-n**0.5//1))+1):
        if tmp%i==0:
            cnt=0
            while tmp%i==0:
                cnt+=1
                tmp //= i
            arr.append([i, cnt])
    if tmp!=1:arr.append([tmp, 1])
    if not arr:arr.append([n, 1])
    return arr

cnt_d = defaultdict(int) # それまでに出てきた「平方数で割った余りの数」の個数
ans = 0

for ｊ,a in enumerate(A_L):
    if a == 0:
        ans += j # j := aより左側にある数字の数
        cnt_d[0] += 1
        continue
    tmp = 1
    for p,p_cnt in factorization(a): # 素因数分解して、指数が偶数の素数を掛けていく
        if p_cnt % 2 == 0:
            continue
        tmp *= p

    ans += cnt_d[tmp] + cnt_d[0] # 平方数で割った余りの数 + 0の数を追加
    cnt_d[tmp] += 1
print(ans)