AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC136 E - Max GCD

問題原文

atcoder.jp

問題要旨

長さ  N の整数列  A がある。「ある数に +1し、ある数に-1する」という操作を  K 回以下行ってできる  A' の全ての数を割り切るような数のうち 、最大のものは何か?ただし操作の途中で要素の値が負になってもよい。

  •   2 \leq N \leq 500
  •   1 \leq A \leq 10^{6}
  •   0 \leq K \leq 10^{9}

解法

答えの候補を絞る

操作後に  X で全ての  a'_i を割りきることができるということは、 sum(A') X で割り切れる。
また、操作を行っても  sum(A) の和は一定であることに注目すると、  sum(A') を割り切るような  X 以外は答えになり得ないことがわかる。
つまり答えの候補は、  sum(A) の約数である。

操作をして達成できるかを考える

次に、「操作後の  a'_i を全て割り切るような  X を設定した時に、  K 回の操作で 各  a_i を 目的の a'_iにすることができるか?」を考える。
少し考えると、これは「上に合わせるか」「下に合わせるか」の二択があることがわかる。
例えば、  a_i = 10,  X=7 だとして、 10を7で割り切れるような数にするには、 「上に合わせて10 + 4 = 14 にする」か、「下に合わせて 10 - 3 = 7」にするかということである。

例えばサンプル3の A = 1, 2, 10 ,22 という例で、  X = 7 として考えるとこんな感じ。
まずは、  X で modをとってソートしておくと

  • A = 1, 1, 2, 3

になる。 この上で、上( X)に合わせる場合と下(0)に合わせる場合を考えると、

  • U = +6, +6, +5, +4 (上に合わせる場合)
  • D = -1, -1, -2, -3(下に合わせる場合)

のようになる。
ここから、 - 採用した操作について、回数の和が  K 回以下
- 上に合わせるコストの総和 = 下に合わせるコスト総和

を満たすように、各項について上に合わせるか下に合わせるか決めていく。

そして実は「ある項までは下に合わせて、それ以降の項は上に合わせる」とするのが最適。
たとえば、  A = 1, 1, 3, 5,  X = 7 として、↓↓↑↓のように合わせるとすると、上に合わせるコストの総和は4で、下に合わせるコストの総和は7になる。この時、+1と-1の辻褄のずれは3。全体の操作回数は4回になる。
だが3番目と4番目の項の合わせる向きを逆にして、↓↓↓↑のようにすると、上に合わせるコストの総和 は2で、下に合わせるコストの総和は5になる。このとき、+1と-1の辻褄のずれは3と変わらないが、全体の操作回数を明らかに少なくできる。
よってある区切りを設けて上と下のどちらに合わせるかを決めて良い。また明らかに操作回数が少なくなるように小さい方からみていくのがいい。
実際の判定は累積和を用いると高速にできる。

条件を満たすような区切り方がある  X の MAXが答えになる。

実装

from itertools import accumulate
N, K = map(int, input().split())
A = list(map(int, input().split()))
M = sum(A)
 
# 答えの候補を列挙
ans_candidates = []
for n in range(1, int(M ** 0.5) + 1):
    if M % n == 0:
        ans_candidates.append(n)
        ans_candidates.append(M // n)
 
 
ans = 0
for X in ans_candidates:
    A_mod = sorted([a % X for a in A])
    U = [X - a for a in A_mod]
    D = [-a for a in A_mod]  # わかりやすいので
 
    U = list(accumulate(U))
    D = list(accumulate(D))
 
    for i in range(N):
        if -D[i] > K:
            break
 
        if -D[i] == (U[-1] - U[i]):
            ans = max(ans, X)
            break
 
print(ans)

感想

難しい。
「操作後の総和が変わらない」「各要素が  X で割り切れるなら、その総和も  X で割り切れる」というのは有用な考え方そう。
あと先に  X で modをとっておくと考えやすさが段違いになる。