何度かTLEしながらなんとか答えにたどり着いたので、考え方をメモしておきます。
使用する言語はPythonです。
問題文
N個の整数 A1,A2,…,ANが黒板に書かれています。
あなたはこの中から整数を1つ選んで、1以上109以下の好きな整数に書き換えます。
元の整数と同じ整数に書き換えても構いません。
書き換えた後のN個の整数の最大公約数の最大値を求めてください。制約
入力は全て整数である。
2≤N≤105
1≤Ai≤109
最初に書いたプログラム
import fractions from functools import reduce def gcd_list(numbers): return reduce(fractions.gcd, numbers) n = int(input()) a = list(map(int,input().split())) ans = 0 for i in range(n): li = a[:i] + a[i + 1:] tmp = gcd_list(li) if tmp > ans: ans = tmp print(ans)
“整数を1つ選んで、1以上109以下の好きな整数に書き換える”とは、”整数を1つ選んで取り除く”ことと同じであると解釈できます。
なので、a[:i] + a[i + 1:]で整数を1つ取り除いたリストの全パターンを作り、最大公約数を毎回計算しました。
しかし、N-1個の整数の最大公約数を取る処理をN回行うのでTLEしました。
累積和を応用した累積GCDという考え方を使えば、大幅に計算量を減らすことができるようです。
累積GCDを使って計算量を減らす
仮に、入力を
64 48 16 18 12
だった場合を考えてみましょう。
僕が最初に書いたプログラムは、
1回目のループ
[48, 16, 18, 12]から最大公約数を求める
2回目のループ
[64, 16, 18, 12]から最大公約数を求める
といったように進んでいきます。
このとき、右側の3つの数字[16, 18, 12]の最大公約数を2回求める必要があります。
[16, 18, 12]の最大公約数を保持しておけば、48や64との最大公約数を求めるだけで計算できます。こういった、重複する処理を削減するために累積GCDを使います。
そもそも累積和を忘れてしまった方は、こちらの記事が非常に参考になります。
復習しておきましょう。
左側からの累積GCDを以下のように用意します。
- left[0] = 0
- left[1] = A0の最大公約数(要するにA0)
- left[2] = A0とA1の最大公約数
- left[3] = A0とA1とA2の最大公約数
- ︙
- left[N + 1] = A0とA1とA2と…とANの最大公約数
また、右側からの累積GCDも用意しておきましょう。
(最初のプログラムでいうa[i + 1:]の部分が必要なため)
import fractions from functools import reduce n = int(input()) a = list(map(int,input().split())) #累積GCD left = [0] right = [0] for i in range(n): left.append(fractions.gcd(left[i],a[i])) right.insert(0,fractions.gcd(right[0],a[n - 1 - i])) print(left) print(right) ans = 0 for i in range(n): ans = max(ans,fractions.gcd(left[i],right[i + 1])) print(ans)
このコードを実行すると、
PS D:\Atcoder> python c.py [0, 64, 16, 16, 2, 2] [2, 2, 2, 6, 12, 0] 4
と出力されます。
わかりやすいように、左右からの累積GCDも出力しました。
累積GCDでもTLEする原因と対策
累積GCDによって計算量はかなり減少しました。
しかし、このコードではまだTLEしてしまいます。
累積GCDを求める際に、appendで一つずつリストに追加していることが原因です。
appendを使うと、リストの中身をコピーして作り直します。
競プロでは実行時間を極力短くしたいので、使うのは避けたほうが良さそうです。
TLEするのは先頭へのinsertがO(n)であることが原因だとコメントをいただきました。
よく使う操作はOを調べるよう注意することにします。
今回は、用意するリストの大きさが決まっているので、
left = [0] * (n + 1)
right = [0] * (n + 1)
と書いて領域を確保しておきます。
import fractions from functools import reduce n = int(input()) a = list(map(int,input().split())) #累積GCD left = [0] * (n + 1) right = [0] * (n + 1) for i in range(n): left[i + 1] = (fractions.gcd(left[i],a[i])) right[n - i - 1] = (fractions.gcd(right[n - i],a[n - 1 - i])) ans = 0 for i in range(n): ans = max(ans,fractions.gcd(left[i],right[i + 1])) print(ans)
これを時間内に閃くって相当難易度高いんじゃないですかね。
Atcoder参戦して2回目の人間にはかなりハードルが高かったです。
早く茶色まで上がりたい。