AtCoder が提供しているABC（AtCoder Beginner Contest）339 のF問題をC++とPythonで解いてみました。ABC339は、2024年2月3日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
F問題 Product Equality（Difficulty : 1716）
問題はリンク先をご覧ください。
大きな素数のMOD付の演算で判断します。AtCoder Problems による Difficulty は 1716 でした。
解答案
$1 \leqq i, j, k \leqq N\, (1 \leqq N \leqq 1000) $ となる $i, j, k$ に対して、$A_i \times A_j = A_k$ となる $(i, j, k)$ の組の数を求める問題です。$N$ が小さいこともあり、それほど難しくない問題と思えます。ただし、
$$1 \leqq A_i \leqq 10^{1000}$$
という制約があります。
もし、$A_i \times A_j = A_k$ が成り立っていれば、$A_i, A_j, A_k$ の剰余（mod）をとっても式は成立します。ただし逆は成立しません。
ある程度大きな素数を選んで、その剰余付き計算で、$A_i \times A_j = A_k$ となる $(i, j, k)$ の組の数を求めてみます。
これは、確率的に間違う可能性があります。紹介したプログラムでは、2つの素数を選ぶことによって、ACを取ることができました。
C++ プログラム例（ABC339F）
プログラムを補足します。
- 10億から20億の間の素数2つを適当に選びます（6、7行目）。
- 32ビット整数は、符号有りで最大約21億、符号無しで最大約42億です。符号有りの最大値に近い素数を2つ選びました。大きな素数は検索して見つけました。
- 文字列を整数と考えて、剰余を計算する関数 convert を用意します（9ー20行目）。
- 素数1、素数2で剰余を取った配列を c1、c2 とします。それぞれの配列の要素別の個数を map コンテナ m1、m2 に格納します（35ー40行目）。
- p? = c?[i] × c?[j] を計算して、コンテナからその要素の個数を求めて、結果に加えます。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
#define MOD1 1099000027
#define MOD2 1999000021
ull convert(string s, ull mod)
{
ull result = 0;
for (int i = 0; i < s.length(); ++i) {
result *= 10;
result += s[i] - '0';
result %= mod;
}
return result;
}
int main()
{
int n;
cin >> n;
vector<string> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<ull> c1(n);
map<ull, int> m1;
vector<ull> c2(n);
map<ull, int> m2;
for (int i = 0; i < n; ++i) {
c1[i] = convert(a[i], MOD1);
++m1[c1[i]];
c2[i] = convert(a[i], MOD2);
++m2[c2[i]];
}
int result = 0;
for (int i = 0; i < n; ++i) {
for (int j = i; j < n; ++j) {
ull p1 = (c1[i] * c1[j]) % MOD1;
ull p2 = (c2[i] * c2[j]) % MOD2;
int t = min(m1[p1], m2[p2]);
if (i == j) {
result += t;
} else {
result += 2 * t;
}
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC339F）
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 339 F"""
from collections import defaultdict
MOD1 = 1099000027
MOD2 = 1999000021
def convert(s, mod):
result = 0
for i, ch in enumerate(s):
result *= 10
result += ord(ch) - ord('0')
result %= mod
return result
n = int(input())
a = [""] * n
for i in range(n):
a[i] = input()
c1 = [0] * n
m1 = defaultdict(int)
c2 = [0] * n
m2 = defaultdict(int)
for i in range(n):
c1[i] = convert(a[i], MOD1)
m1[c1[i]] += 1
c2[i] = convert(a[i], MOD2)
m2[c2[i]] += 1
result = 0
for i in range(n):
for j in range(i, n):
p1 = (c1[i] * c1[j]) % MOD1
p2 = (c2[i] * c2[j]) % MOD2
t = min(m1[p1], m2[p2])
if i == j:
result += t
else:
result += 2 * t
print(result)Python は、整数の大きさに上限はありません。以下の素直な実装が Python (PyPy 3.10-v7.3.12) で AC しました。ただし、これは想定解ではなく、時間もぎりぎり間に合っているだけです。
"""AtCoder Beginner Contest 339 F"""
from collections import defaultdict
n = int(input())
a = [0] * n
for i in range(n):
a[i] = int(input())
c = defaultdict(int)
for i in range(n):
c[a[i]] += 1
result = 0
for i in range(n):
for j in range(i, n):
p = a[i] * a[j]
if p in c:
if i == j:
result += c[p]
else:
result += 2 * c[p]
print(result)こちらも「AC」と判定されました。
最後に
与えられる整数の桁数が大きいため、コンテスト中は、Python を使って考えていました。公式解説をみて仕組みが理解できました。
引き続き ABC の問題を紹介していきます。
