累積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

最初に書いたプログラム

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:]で整数を１つ取り除いたリストの全パターンを作り、最大公約数を毎回計算しました。

しかし、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:]の部分が必要なため）

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回目の人間にはかなりハードルが高かったです。

室井

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