公式

A - 階段の水流 / Waterflow on Stairs 解説 by admin

GPT 5.2 High

概要

各操作で指定された植木鉢の水を「\(B_i\) だけ蒸発させ、残りを1つ下へ流す」ことを順にシミュレーションし、最終的な各植木鉢の水量を求めます。

考察

この現象で重要なのは、1回の日光操作で水が移動するのは高々1段だけという点です。
つまり、\(i\) 番目に日光を当てたら

  • \(i\) 番目の水は必ず \(0\) になる(蒸発+残りは全流出)
  • 流れるのは蒸発後の残り \( \max(0, A_i - B_i) \) だけ
  • その残りは \(i+1\) 番目に足されるだけ(\(i=N\) なら排水溝に消える)

となります。流れ落ちた先(\(i+1\) 番目)でさらに蒸発が起きるのは、将来その鉢が操作対象に選ばれたときだけです。

素朴に「操作のたびに全ての鉢の状態を見直す」ような実装をすると \(O(NQ)\) となり、制約 \(N,Q \le 2\times 10^5\) では間に合いません。
しかし実際には、各操作で影響を受けるのは 鉢 \(i\) と(存在すれば)鉢 \(i+1\) のみなので、1回あたり \(O(1)\) で更新できます。

アルゴリズム

現在の水量配列を w[0..N-1] として持ち、操作ごとに次を行います(1-index を 0-index に直して扱う):

  1. 操作対象を \(i\) とする(0-index)
  2. 現在の水 cur = w[i] を取り出す
  3. 蒸発後に流れる量は out = max(0, cur - B[i])
  4. w[i] = 0 にする(残りは全て流れるため)
  5. \(i+1 < N\) なら w[i+1] += out(最下段なら捨てる)

例えば \(N=3\)、水 \((10,0,0)\)、\(B=(3,5,0)\) で鉢1に日光を当てると、 蒸発で \(3\) 消え、残り \(7\) が鉢2へ移るので \((0,7,0)\) になります。

これを \(Q\) 回繰り返せば最終状態が得られます。

計算量

  • 時間計算量: \(O(N+Q)\)(初期読み込み+各操作 \(O(1)\))
  • 空間計算量: \(O(N)\)(水量配列)

実装のポイント

  • 操作では必ず w[i]=0 になる点を見落とさない(「残りを下へ全部流す」ため)。

  • 最下段 \(i=N\) のときは排水溝へ消えるので、i+1 への加算は行わない。

  • \(A_i, B_i\) は最大 \(10^9\)、加算で大きくなり得るため、十分大きい整数型(Python の int)で扱う。

    ソースコード

import sys


def read_all_ints():
    data = sys.stdin.buffer.read()
    res = []
    num = 0
    in_num = False
    for b in data:
        if 48 <= b <= 57:
            num = num * 10 + (b - 48)
            in_num = True
        else:
            if in_num:
                res.append(num)
                num = 0
                in_num = False
    if in_num:
        res.append(num)
    return res


def main():
    v = read_all_ints()
    p = 0
    N = v[p]
    Q = v[p + 1]
    p += 2

    water = v[p:p + N]
    p += N
    B = v[p:p + N]
    p += N

    w = water
    b = B
    n = N
    for _ in range(Q):
        i = v[p] - 1
        p += 1
        cur = w[i]
        bi = b[i]
        if cur > bi:
            out = cur - bi
            if i + 1 < n:
                w[i + 1] += out
        w[i] = 0

    sys.stdout.write(" ".join(map(str, w)))


if __name__ == "__main__":
    main()

この解説は gpt-5.2-high によって生成されました。