目次
はじめに
この記事では、ABC342のD問題を解説していきます。
ABC（AtCoder Beginner Contest）とは、AtCoderが開催している、競技プログラミングコンテストです。
ABC342 D – Square Pair
問題
長さNの非負整数列 \(A=(A_1, …, A_N)\) が与えられます。積が平方数となる2数のインデックスの組の個数を数える問題です。
制約
- \(2 \leqq N \leqq 2×10^5 \)
- \(0 \leqq A_i \leqq 2×10^5 \)
思考の筋道
この問題において、0は特別な数です。0はもう片方の数が何であっても、積が0で平方数になるからですね。そのため、まずは0以外の数について考えてみましょう。
0以外の数の検討
0以外の数については素因数分解を行います。2つの数の積が平方数になるのは、積の素因数の指数が偶数になるときです。つまり、かける前の2つの数について、指数が奇数の素因数が完全に一致するときですね。少しややこしいので例を見てみましょう。
具体例：2と72
なお、2つの数について、指数が奇数の素因数が完全に一致するかどうかは、指数が奇数の素因数の積が一致するかどうかで調べることができます。
0を含む場合の検討
次は0を含むケースを考えましょう。はじめに述べた通り、0はもう片方の数が何であっても、積が0で平方数になります。しかし、0が複数ある場合は、0×0をダブルカウントしないように注意が必要です。
具体例：A=(0, 0, 2, 4, 8, 72)
- 1つ目の0について
0はもう片方の数が何であっても平方数になるので、自分を除いた5だけカウントプラス - 2つ目の0について
先ほど「1つ目の0×2つ目の0」のケースは数えたので、そのケースと自分を除いた4だけカウントプラス - 残りの2, 4, 8, 72について
この中で指数が奇数の素因数が完全に一致するのは、2, 8, 72の三つ。この中から二つを選べば積が平方数になるので、\({}_3 \mathrm{ C }_2=3 \) だけカウントプラス
よって答えは、5+4+3=12となる。
整数問題で素因数に着目する問題はしばしば出題されます。
下記のコードでは、factorization関数で素因数分解をしています。素因数分解のコードはいつでも取り出して使えるようにしておくとよいでしょう。
コード
N = int(input())
A = list(map(int, input().split()))
def factorization(n):
arr = []
temp = n
for i in range(2, int(-(-n ** 0.5 // 1)) + 1):
if temp % i == 0:
cnt = 0
while temp % i == 0:
cnt += 1
temp //= i
arr.append([i, cnt])
if temp != 1:
arr.append([temp, 1])
return arr
def f(a):
fac = factorization(a)
res = 1
for k, n in fac:
if n % 2:
res *= k
return res
d = dict()
ans = 0
zero_count = 0
for num in A:
if num == 0:
ans += N - 1 - zero_count
zero_count += 1
else:
s = f(num)
if s in d:
d[s] += 1
else:
d[s] = 1
for a in d:
if d[a] >= 2:
ans += (d[a] * (d[a] - 1)) // 2
print(ans)以上、ABC342のD問題の解説でした！
では、またね。
2は、指数が奇数の素因数は2のみ。
72(=2^3×3^2)も、指数が奇数の素因数は2のみ。
この2つの数は、指数が奇数の素因数が一致します。かけると 2^4×3^2 となり、すべての素因数の指数が偶数となります。そのため、平方数です。