ABC-393 E - GCD of Subset
https://atcoder.jp/contests/abc393/tasks/abc393_e
コンテスト中には解けませんでしたがコンテスト後に調べて正解できました。その過程でいろいろなことがわかったので、長くなりますがまとめます。簡単にいうといろんなアプローチを試して結局無理で、最終的には公式解説をなぞる形になりました。
コンテスト中の自分の動き
D問題まではすんなりと解くことができ、60分以上残っていましたがE, F問題は全くわかりませんでした。このE問題に関してはとりあえず約数ということから素因数分解をしてみて、素因数を眺めつつ何かひらめかないかなと考え込みましたが何もでてきませんでした。計算量については N <= 1.2 * 10^6 よりも A[i] <= 10^6 を利用しないといけないのかなと思い、先に 10^6 までの整数を全て素因数分解しておくことを考えました。しかし素因数分解には O(√N) の計算量がかかるはずなので O(N√N) となり、10^9 回ぐらいの計算回数になってTLEだな……などと考えていました。
誤ったアプローチ（約数全列挙）
コンテスト終了後に Twitterでいろんな人のツイートを見ていたところ、素因数分解ではなく約数に着目すべきということがわかりました。「各 A[i] についてそれぞれの約数を列挙しておく。K 個以上の A[i] がその約数を持っているとき、その約数は A から K 個選んだときの公約数になり得る。」とのことです。素因数分解が思いつくのなら約数の列挙ぐらいはすぐに思いつくべきでしたね。頭が固いです。
しかし約数の列挙にはやはり O(√N) の計算量が必要になります（1 から √N までの整数で割りきれるかどうかを調べるやり方）。これを 10^6 まで繰り返したら TLE になります。いろいろ調べてみたら素因数分解を高速化する方法（osa_k法）が紹介されていました。
先に改造版のエラトステネスのふるいを行ってそれぞれの整数に対しそれを割り切れる最小の素数を記録しておき、その情報を利用して高速に素因数分解を行う方法です。これは O(NlogN) とのことなので時間内に間に合いそうです。この方法で素因数分解をし、それを利用して約数を列挙します。それぞれの素因数を掛け合わせるパターンを全て試すやり方を使います。約数の個数を調べるときによく用いられる考え方ですね。イメージとしては例えば 2^3 x 3^2 x 5^1 = 360 の場合、2 は 0 乗から 3乗まで、3は0乗から2乗まで、5は0乗から1乗までなので 4種類x3種類x2種類 = 24通りみたいになります。これを全探索すると約数の総数の分だけ計算することになりますが、間に合うことを期待したいです。実際、下記のコードを用いて10^6までの整数について計算してみたら 13,970,035 ≒ 1.4 * 10^7 個でした。
頑張って バックトラックを利用した DFS を書いてみました。長くなりますがせっかく書いたので一応コードを貼っておきます。10^6までの素数を求める→10^6までの整数を全て素因数分解する→DFSで約数を全て列挙するという流れです。ちなみにAtCoderのコードテストではPyPy3を使用したところ 3539ms を要しました。遅いです。DFSで何度も関数を呼び出しているせいかもしれません。
# 1 から max_num までの整数について、その約数を列挙する
def enumerate_divisors(max_num): # 改造版エラトステネスのふるい
# max_num 以下の素数を得る + ある整数に対しそれを割り切れる最小の素数を記録する
def eratosthenes(max_num):
is_prime = [1 for _ in range(max_num + 1)] # 1 なら素数、それ以外なら素数でない
is_prime[0], is_prime[1] = 0, 0
for p in range(2, max_num+1):
if is_prime[p] != 1: continue # p が既に合成数と判定されていればスキップ
q = p * 2
while(q <= max_num):
if is_prime[q] == 1:
is_prime[q] = p
else:
is_prime[q] = min(p, is_prime[q])
q += p
return is_prime
# 1 以上 max_num 以下の整数を全て素因数分解する
is_prime = eratosthenes(max_num)
def prime_factorize(num):
if is_prime[num] == 1:
return [(num, 1)]
else:
prime_factors = []
for i in range(is_prime[num], num + 1):
if i * i > num: # √max_num までを調べる。これを2乗で表現する。
break
if num % i == 0: # num を i で何回割ることができるか
j = 0
while(num % i == 0):
num //= i
j += 1
prime_factors.append((i, j))
if num != 1: # 素数が最後に残ることがありうる
prime_factors.append((num, 1))
return prime_factors
prime_factors = [[(0, 0)], [(1, 1)]]
for i in range(2, max_num + 1):
prime_factors.append(prime_factorize(i))
# 整数 num の約数を列挙する
divisors = [set() for _ in range(max_num + 1)]
def enumerate_factors(num):
def dfs(x, index, prime_factors): # DFSバックトラック
if index >= len(prime_factors):
divisors[num].add(x)
else:
[n, multi] = prime_factors[index]
for m in range(multi + 1):
x *= n ** m
dfs(x, index + 1, prime_factors)
x //= n ** m # バックトラック
dfs(1, 0, prime_factors[num])
for i in range(1, max_num + 1):
enumerate_factors(i)
return divisors
enumerate_divisors(10**6)
もっと速い約数全列挙
さらに調べてみたところ、さらに速い方法が見つかりました。これもどなたか（すみません、誰だったか忘れました）がTwitterで書かれていたアイデアから来ています。感覚的にはエラトステネスのふるいと同じですね。ある整数 x について、x で割り切れる数字全てについて x を記録しておく感じです。結局のところアイデアとしては先に書いた O(NlogN) の素因数分解と同じです。
def enumerate_divisors2(max_num):
divisors = [[1] for _ in range(max_num + 1)]
for p in range(2, max_num + 1):
q = p
while (q <= max_num):
divisors[q].append(p)
q += p
return divisors
enumerate_divisors2(10**6)
これはコードテストで 1285ms となりました。これなら間に合いそうです。コードもシンプルでとても良いですね。
しかしこれを利用して解いてみたところ、やはり TLE でした。おそらくですが「A に含まれる全ての約数」を走査してしまったのだと思います。A[i] の中に約数の数が多い整数が大量に入っていたら遅くなる気がします。時間内に間に合うのはあくまで「10^6までの整数に含まれる全ての約数」の走査です。
ちなみにダメ元で ChatGPT に C++ のコードに変換してもらったらなんと AC になりました。こういうことが増えてくるようなら C++ への乗り換えをいつかは考えないといけないのかもしれません。
考察
ここでようやく、複雑そうだからなんとなく敬遠していた公式解説を丁寧に読んで、なんとかその方針に従って AC できました。答えからの逆算にはなりますが、自分なりにこんな風に考えていけば解けるかな？という流れを書いてみます。読みづらいかもしれません……。
愚直に考えるとN個からK個選んでgcdを取ればいいわけですが、もちろんTLEなのでやりません。そこで各A[i]が持つ約数をあらかじめ列挙しておき、その中で数列Aの中にK個以上存在する約数の一覧を取得しておき、各A[i]が持つ約数の中で最大のものを答えとすることを考えます。先に書いた「誤ったアプローチ」ですね。これもTLEになるので次の方法を考えます。
約数ではなく、逆に倍数の方から考えてみます。「約数ではなく倍数から考える」といえばエラトステネスのふるいです。これは調和級数の性質（※）を利用して計算量を小さくできるやり方です。
（※）調和級数の計算量
この級数は「各項の値は 0 という極限を持つにも関わらず、その級数の和は発散する」という性質を持っています。これは音楽理論などにおける倍音の概念に由来して、調和級数(Harmonic Series)と呼ばれています。
とのことです。また、
https://manabitimes.jp/math/627#4
より 1/1 + 1/2 + 1/3 + ...... + 1/N は log(N+1) より大きく、logN + 1 より小さいそうです。計算量を考えるときはほぼ logN と考えていいと思います。
上の式がここでいう「劣化エラトステネスのふるい」になります。
https://qiita.com/drken/items/3beb679e54266f20ab63
正式版は分母に来る数字が素数に限られるため、計算量がもっと小さくなるようです。
コードにするとこうですね。
for x in range(1, MAX+1):
y = x
while (y <= MAX): # このループ回数が MAX * (1/1 + 1/2 + 1/3 + ...... ) に収まる
y += x
考察の続き
ここで先に書いた osa_k 法を思い出します。この方法ではエラトステネスのふるいを改造して、素数か否かの判定だけでなく「ある整数に対しそれを割り切れる最小の素数を記録する」という関数を作りました。今回の問題で欲しいのは最小ではなく最大の数（素数とは限らない）です。osa_k 法では僕が前節で書いた enumerate_divisors(max_num) 関数のように、各数字について「2で割り切れる」「3で割り切れる」という情報を記録し、最小の数が入るようにします。以下に A[i] <= 30 の場合を書いてみました。
これを改造して最小ではなく最大の数が入るようにすると下の図のように1から30までの値がそのまま入ります。当たり前ですね。
今回の問題でいうと K=1 のときにこうなります。イメージとしてはここから K を大きくするにつれてこれらの数値が小さくなっていくわけですね。ここまで考えると、じゃあ 1 <= x <= 30 となる x の中で、x の倍数が K 個以上あるものはどれなのかを把握したくなります。となれば、1から30までの数字について、数列 A の中に x の倍数が何個あるのかを数える必要が出てきます。
ここからは問題の制約（A[i] <= 10^6）に沿って考えていきます。公式解説のようにまず A の中にどの数字が何個あるのかを記録しておき、それを見ながら 2 の倍数は何個ある、3の倍数は何個ある……という計算を 10^6 の倍数まで続けます。
# A の中に x の倍数がいくつあるかを num_of_multi_of[x] に格納する。
num_of_multi_of = [0 for _ in range(MAX+1)]
count_A = [0 for _ in range(MAX+1)]
for a in A:
count_A[a] += 1
for x in range(1, MAX+1):
y = x
while (y <= MAX): # ここのループ回数が log(MAX) 未満に収まる（※）
# x の倍数 y が A に含まれていれば、その個数分カウントする。
num_of_multi_of[x] += count_A[y]
y += x
# ※ 調和級数。Σ(1/k) < log N + 1
それが記録できたら次は改造版のエラトステネスのふるいのように、x の倍数が A の中に K 個以上あるような x それぞれを見て、1 <= y <= 10^6 となる全ての整数について「これは x で割り切れる」という情報を付与していきます。maxを取ることで最大値が記録できます。すなわち個々の整数に対し問題の条件を満たす最大の約数が記録できます。
# x の倍数 y （= x は y の約数）について、ans[y] に x を格納していく。
# ただし x の倍数が A の中に K 個以上存在していないといけない。
ans = [0 for _ in range(MAX+1)]
for x in range(1, MAX+1):
if num_of_multi_of[x] < K: continue
y = x
while (y <= MAX): # ここのループ回数が log(MAX) 未満に収まる（※）
ans[y] = max(ans[y], x) # 最大公約数がほしいので max をとる
y += x
ここまでできたらあとは O(N) で答えが求められますね。
実装
N, K = map(int, input().split())
A = list(map(int, input().split()))
MAX = 10**6
# A の中に x の倍数がいくつあるかを num_of_multi_of[x] に格納する。
num_of_multi_of = [0 for _ in range(MAX+1)]
count_A = [0 for _ in range(MAX+1)]
for a in A:
count_A[a] += 1
for x in range(1, MAX+1):
y = x
while (y <= MAX): # ここのループ回数が log(MAX) 未満に収まる（※）
# x の倍数 y が A に含まれていれば、その個数分カウントする。
num_of_multi_of[x] += count_A[y]
y += x
# ※ 調和級数。Σ(1/k) < log N + 1
# x の倍数 y （= x は y の約数）について、ans[y] に x を格納していく。
# ただし x の倍数が A の中に K 個以上存在していないといけない。
ans = [0 for _ in range(MAX+1)]
for x in range(1, MAX+1):
if num_of_multi_of[x] < K: continue
y = x
while (y <= MAX): # ここのループ回数が log(MAX) 未満に収まる（※）
ans[y] = max(ans[y], x) # 最大公約数がほしいので max をとる
y += x
for a in A: # 答えの出力
print(ans[a])
おわりに
振り返ってみると「エラトステネスのふるいをどのぐらい熟知していますか？」という問題だったように思います。これまでは漠然と「素数を速く列挙できる方法」ぐらいにしか捉えていなかったのですが、この問題を通して「約数ではなく倍数から考える」「調和級数の性質により計算量が小さくなる」ということが学べました。