AtCoder が提供しているABC（AtCoder Beginner Contest）342 のD問題をC++とPythonで解いてみました。ABC342は、2024年2月24日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Square Pair（Difficulty : 944）
問題はリンク先をご覧ください。
与えれた整数を平方数で割れるだけ割った値を求めて処理します。AtCoder Problems による Difficulty は 944 でした。
解答案
すべての組み合わせの掛け算の結果を求めて平方数か判断しては、$O(N^2)$ となり間に合いません。
平方数は、素因数分解したときにすべての素因数のべき乗が偶数になります。素因数の奇数のべき乗を持つ数に別の数字を掛けて平方数にするためには、同じ素因数の奇数のべき乗をもつ必要があります。
入力例2の場合、与えられた数は以下です。
2 2 4 6 3 100 100 25
この数を平方数で割れるだけ割ると以下となります。
2 2 1(=4/22) 6 3 1(=100/102) 1(=100/102) 1=(25/52)
この変換で平方数は1となります。平方数4個から2個を選ぶ組み合わせは、$_4C_2 = 6$ 通り、同じ残りを持つ数は、2が2つあるだけです。この組み合わせは1通りですので、合わせて、7通りが解となります。
C++ プログラム例（ABC342D）
考察に従い、与えられた整数を平方数で割れるだけ割った余りを求める関数 remain を実装します。
- 戻り値 result の初期値を1とします。
- 2以上の数で割れるだけ割った回数を求めます。割った回数が奇数の場合、その素因数を result に掛けます。
- 割り切った残りの数は、result に掛けて、result を戻します。
本体 main では、0 の値をカウントします。0 でない場合、関数 remain の戻り値を map コンテナ m に格納します。
最終的は解は、m のそれぞれの要素から求めます。
- mの値 num は、その素因素の余りを持つ数の個数となるため、$_{num}C_2$ がその数を掛けて平方数が得られる組み合わせの数となります（49行目）。
- 値0の場合は、どの数字を掛けても値0となり平方数となります。0が複数ある場合も考慮して、別に計算します（46行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
ll remain(ll n)
{
ll result = 1;
for (ll i = 2; i * i <= n; ++i) {
ll count = 0;
while (n % i == 0) {
++count;
n /= i;
}
if (count % 2 == 1) {
result *= i;
}
}
if (n != 1) {
result *= n;
}
return result;
}
int main()
{
int n;
cin >> n;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
ll zero = 0;
map<ll, ll> m;
for (int i = 0; i < n; ++i) {
if (a[i] == 0) {
++zero;
} else {
++m[remain(a[i])];
}
}
ll result = zero * (n - 1) - zero * (zero - 1) / 2;
for (auto e : m) {
ll num = e.second;
result += num * (num - 1) / 2;
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC342D）
Python 版も基本的な考え方は同じです。与えれた整数を平方数で割った残りは、defaultdict に格納しました。以下となります。
"""AtCoder Beginner Contest 342 D"""
from collections import defaultdict
def remain(n):
result = 1
i = 2
while i * i <= n:
count = 0
while n % i == 0:
count += 1
n = n // i
if count % 2 == 1:
result *= i
i += 1
if n != 1:
result *= n
return result
n = int(input())
a = list(map(int, input().split()))
zero = 0
m = defaultdict(int)
for i in range(n):
if a[i] == 0:
zero += 1
else:
m[remain(a[i])] += 1
result = zero * (n - 1) - zero * (zero - 1) // 2
for num in m.values():
result += num * (num - 1) // 2
print(result)こちらも「AC」と判定されました。
最後に
この問題は、素朴な方法では時間が間に合わないことが分かりました。奇数のべき乗の素因数を持つアイデアから、なんとか時間内に AC 判定を取ることができました。自分にとっては、難易度が高い問題で、解くことができて良かったです。
引き続き ABC の問題を紹介していきます。
