ABC136 E問題の解法(Python)

復習によってABC136のE問題の解法を理解したので、ご紹介します。

E - Max GCD

まず、この問題を解く上での一番の肝は「数の集合Aについて、のすべての要素を正の整数Xで割り切れる」=「数の集合Aの総和が正の整数Xで割り切れる」ということに気付けるかどうかです。
この問題は、何回操作を繰り返してもAの総和が変わらないので、sum(A)の約数をすべて試していけば良さそうです。
約数のリストはO(√N)の計算量で以下のように作成できます。

N, K = map(int, input().split())
A = list(map(int, input().split()))

sumA = sum(A)
divisors = []
for i in range(1, sumA + 1):
    if i * i > sumA:
        break
    if sumA % i != 0:
        continue
    divisors.append(i)
    if sumA // i != i:
        divisors.append(sumA // i)

次は、ここで求めた約数を降順にソートしたリストの要素Xに対して、K回以内の操作で条件を満たせるかどうかを試していきます。
各要素プラス方向もしくはマイナス方向に調整をしてすべての要素がXを約数に持つように調整します。
例えば、A=[8,20]の場合、Aの総和の約数は[1,2,4,7,28]となります。
すべての要素が7を約数に持つようにしたければ、

[8−1,20+1]と調整する(操作に必要な回数=1回)
[8+6,20−6]と調整する(操作に必要な回数=6回)

の2パターンがあります。
以上のように、プラス方向に調整する変更量の総和sumPとマイナス方向に調整する変更量の総和sumMが等しいとき、各要素がXを約数に持つにはsumP回の操作が必要だということになります。
Aのすべての要素がXを約数に持つように調整するのに必要な最小の回数を求めるには、まず、全てをマイナス方向に調整した場合を考え、プラス方向に変更した方が変更量が少なくなるものから(a%Xが大きいものから)プラス方向の調整に切り替えていきます。
そのような順番で切り替えていくことで、sumPとsumMを小さくすることができるので、最小の操作回数を求めることができます。
そして、sumP==sumMとなったときのsumPがK以下であるとき、K回以下の操作でAの全ての要素について、Xを約数に持たせることができます。
これを踏まえて、先程のコードの続きを書くとこんな感じになります。

# 約数が大きい順に調べる
divisors.sort(reverse=True)

for x in divisors:
    costs = [a % x for a in A]
    sumM = sum(costs)
    sumP = 0

    # a % xが大きい順にプラス方向の調整に変えていく
    costs.sort(reverse=True)
    for j in range(N):
        if sumM == sumP:
            break
        else:
            sumM -= costs[j]
            sumP += x - costs[j]

    # sumP(必要な回数)がK以下に抑えられたら、それが答え
    if sumP <= K:
        print(x)
        break

上記は尺取法を用いた解法ですが、実は、これにはもう一個別解があります。
マイナス方向に調整する要素と、プラス方向に調整する要素の境界を一発で求める方法です。
根幹の考え方は同じなのですが、以下の考え方も踏まえると上記の解法よりも楽に解くことができます。

全部マイナス方向に調整する場合の変更量の配列costsを求め、降順にソートしておく
- このとき、sum(costs)は常にXの倍数となる
- また、sumM−sumP=sum(costs)
どれか一つの調整方向の符号を逆転させると、sumM−sumPがXだけ減る
- costsの大きい順に符号を反転させると、sum(costs)//Xが境目になる
- A=[8,20], X=7, costs=[1,6]の場合
  - 全てマイナス方向に調整した場合の変更量の総和sumM=1+6=7,sumP=0。よって、sumM−sumP=7=sum(costs)
  - 20をプラス方向に調整することにすると、マイナス方向の調整量は6だけ減ってプラス方向の調整量に7−1=6だけ増える。よってsumM−sumP=sum(costs)−7=0
sum(costs[sum(costs)//d:)]が必要な操作回数の最小値。

以上を踏まえるとこんなコードになる。

divisors.sort(reverse=True)

for x in divisors:
    costs = [a % x for a in A]
    costs.sort(reverse=True)
    if sum(costs[sum(costs) // x:]) <= K:
        print(x)
        break

だいぶシンプルになりました。
私は本番では、「数の集合Aについて、Aのすべての要素を正の整数Xで割り切れる」=「数の集合Aの総和が正の整数Xで割り切れる」に気付けなかったので、ここまではたどり着けませんでした。
てか約数を求める方法すら知らなかったので、良い勉強になりました。
ABC137は参加できませんが、これからもちゃんと復習していこうと思います。