累積GCDでAtCoder ABC 125 C – GCD on Blackboardを解く

python Python

何度かTLEしながらなんとか答えにたどり着いたので、考え方をメモしておきます。
使用する言語はPythonです。

問題文

N個の整数 A1,A2,…,ANが黒板に書かれています。
あなたはこの中から整数を1つ選んで、1以上109以下の好きな整数に書き換えます。
元の整数と同じ整数に書き換えても構いません。
書き換えた後のN個の整数の最大公約数の最大値を求めてください。

制約
入力は全て整数である。
2≤N≤105
1≤Ai≤109

C – GCD on Blackboard

最初に書いたプログラム

“整数を1つ選んで、1以上109以下の好きな整数に書き換える”とは、”整数を1つ選んで取り除く”ことと同じであると解釈できます。

なので、a[:i] + a[i + 1:]で整数を1つ取り除いたリストの全パターンを作り、最大公約数を毎回計算しました。

しかし、N-1個の整数の最大公約数を取る処理をN回行うのでTLEしました。

累積和を応用した累積GCDという考え方を使えば、大幅に計算量を減らすことができるようです。

AtcoderのPython3のバージョンが3.4.3のため、mathモジュールではなくfractionsモジュールにgcd()関数があります。

累積GCDを使って計算量を減らす

仮に、入力を

5
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:]の部分が必要なため)

このコードを実行すると、

と出力されます。

わかりやすいように、左右からの累積GCDも出力しました。

累積GCDでもTLEする原因と対策

累積GCDによって計算量はかなり減少しました。

しかし、このコードではまだTLEしてしまいます。
累積GCDを求める際に、appendで一つずつリストに追加していることが原因です。

appendを使うと、リストの中身をコピーして作り直します。
競プロでは実行時間を極力短くしたいので、使うのは避けたほうが良さそうです。

TLEするのは先頭へのinsertがO(n)であることが原因だとコメントをいただきました。
よく使う操作はOを調べるよう注意することにします。

今回は、用意するリストの大きさが決まっているので、

left = [0] * (n + 1)
right = [0]  * (n + 1)

と書いて領域を確保しておきます。

これを時間内に閃くって相当難易度高いんじゃないですかね。
Atcoder参戦して2回目の人間にはかなりハードルが高かったです。

室井

早く茶色まで上がりたい。