ABC136 E問題の解法(Python) - Web系エンジニアのアウトプット練習場

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

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

約数のリストは $O(\sqrt{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

上記は尺取法を用いた解法ですが、実は、これにはもう一個別解があります。

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

全部マイナス方向に調整する場合の変更量の配列を求め、降順にソートしておく
- このとき、 $sum(costs)$ は常に $X$ の倍数となる
- また、 $sumM - sumP = sum(costs)$
どれか一つの調整方向の符号を逆転させると、がだけ減る
- $costs$ の大きい順に符号を反転させると、 $sum(costs) // X$ 番目が境目になる
- ], , ]の場合
  - 全てマイナス方向に調整した場合の変更量の総和 $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は参加できませんが、これからもちゃんと復習していこうと思います。