C - Range Increment Editorial by harurun4635


公式解説とほとんど同じかつ、公式解説のほうが筋がよい気がしてきましたが、一応残しておきます。

以下では \(A_i = 0\) の要素をすべて \(A_i = M\) としておきます。


もっとも簡単に思いつきそうな貪欲は以下のようなものです。

\(i\) の昇順に変数 \(X = 0\) (直前の要素に操作した回数)を管理しながら以下のようにすすむ

  • \(A_i + X \ge M\) のとき

    \(X \gets M - A_i\)
    \(A_i \gets 0\)

  • そうではなく \(A_i + X + K \ge M\) のとき

    \(C = M - (A_i + X)\)
    \(K \gets K - C\)
    \(X \gets X + C\)
    \(A_i \gets 0\)

  • どれでもないとき

    \(X \gets 0\)

この貪欲が嘘なのは、\(1\) つ目のパターンで、実は \(A_i + X = M\) にするのではなく \(A_i + X = 2M\) と揃えるほうがのちのち得するパターンがあるということです。(直感的にも \(X + K\) の総和がそのほうが大きくなるので明らかでしょう)


ということで、それを組み込んだ貪欲を考えます。

\(1\) つ目の遷移のとき、「 コスト \(2M - (A_i + X)\)\(X \gets X + M\) が買える」ということにします。そして、未来で \(X \gets X + M\) が欲しくなったときに「実はこのタイミングで買っていた(そちらの遷移をしていた)」ということにします。

しかし、そのようにすると次に発生する問題が \(2\) つ目の遷移において「過去の \(X \gets X + M\)」と「 \(X \gets X + C\) 」のどちらを買うべきかわからない点です。(コストと得られる \(X\) の組は単純には比較できないからです)

しかし、これも簡単に解決します。後者を「コスト \(C\)\(X \gets X + M\) が買える」としてもよいからです。(これは直感的にはコスト以上のものを買えているためおかしいです)

証明は以下のとおりです。ただし \(X \gets X +M\) を買ったとしても \(A_i + X = M\) の制約から更新は \(X \gets X + C\) であることに注意してください。
もし、このタイミングで後者を買うのであれば \(X \gets X + C\) を結果するため「コスト以上を買っている」事は起きません。
また、このタイミングで前者を買うとします。このとき後者を未来で買いたくなったとき、「このタイミングで実は後者を買い、未来で前者を買った」ことにすればよいです。

よって、すべて \(X \gets X+M\) が売られていますからコストが小さいものを買えばよいです。 \(3\) つ目の遷移の場合も売られているものに \(\le K\) があれば買えることに注意してください。


ということで以上を実装すればよいです。

import sys
from heapq import heappop, heappush
input = sys.stdin.readline

for _ in range(int(input())):
    n, m, k = map(int, input().split())
    a = list(map(int, input().split()))
    
    x = 0
    que = []
    for i in range(n):
        if a[i] == 0: a[i] = m 
        if a[i] + x >= m:
            heappush(que, 2 * m - (a[i] + x))
            x = m - a[i]
            a[i] = 0
        else:
            c = m - (a[i] + x)
            heappush(que, c)
            if que[0] <= k:
                k -= heappop(que)
                x = x + c
                a[i] = 0
            else:
                que.clear()
                x = 0

    print(*a)

posted:
last update: