Official

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

Claude 4.5 Opus

概要

\(N\) 段の階段状に並んだ植木鉢に対して、指定された植木鉢に日光を当てる操作を \(Q\) 回シミュレーションし、最終的な各植木鉢の水の量を求める問題です。

考察

問題の整理

日光を当てる操作は以下の2ステップで行われます: 1. 植木鉢 \(i\) の水が最大 \(B_i\) ミリリットル蒸発する 2. 残った水がすべて植木鉢 \(i+1\) に流れ落ちる(植木鉢 \(i\) は空になる)

重要な気づき

  • 水は上から下へ一方向にのみ流れます
  • 各操作で影響を受けるのは、日光を当てた植木鉢 \(i\) と、その下の植木鉢 \(i+1\) の2つだけです
  • 操作が連鎖的に複数の植木鉢に影響を与えることはありません

素朴なアプローチで十分か?

各操作は定数時間 \(O(1)\) で処理できるため、\(Q\) 回の操作を愚直にシミュレーションしても全体で \(O(Q)\) となります。\(N, Q \leq 2 \times 10^5\) の制約では十分高速です。

アルゴリズム

  1. 配列 water を初期値 \(A\) で初期化する
  2. 各操作 \(c\) について以下を実行する:
    • 植木鉢 \(c\) の現在の水量を取得
    • 蒸発量を計算: \(\min(\text{現在の水量}, B_c)\)
    • 残りの水量を計算: \(\text{現在の水量} - \text{蒸発量}\)
    • 植木鉢 \(c\) の水を \(0\) にする
    • \(c < N\) なら、残りの水を植木鉢 \(c+1\) に加える
  3. 最終的な water 配列を出力する

具体例

\(N=3\), \(A=[10, 5, 3]\), \(B=[3, 2, 1]\) で植木鉢1に日光を当てる場合:

  • 植木鉢1の水: 10 ml
  • 蒸発: \(\min(10, 3) = 3\) ml
  • 残り: \(10 - 3 = 7\) ml が植木鉢2へ流れる
  • 結果: \([0, 12, 3]\)

計算量

  • 時間計算量: \(O(N + Q)\)

    • 初期配列の読み込みに \(O(N)\)
    • 各操作の処理に \(O(1)\) × \(Q\) 回 = \(O(Q)\)
  • 空間計算量: \(O(N)\)

    • 水の量を管理する配列 water に \(O(N)\)

実装のポイント

  1. 0-indexed と 1-indexed の変換: 問題文では植木鉢は1から始まる番号ですが、配列は0から始まるため、i = c - 1 で変換が必要です。

  2. 境界条件の処理: \(i = N\) の場合(最後の植木鉢)、水は排水溝に流れて消えるため、i + 1 < N のチェックを忘れないようにします。

  3. 蒸発量の計算: 水が \(B_i\) 未満の場合はすべて蒸発するため、min(current_water, B[i]) で蒸発量を計算します。

    ソースコード

import sys
from collections import defaultdict

def main():
    input = sys.stdin.readline
    N, Q = map(int, input().split())
    A = list(map(int, input().split()))
    B = list(map(int, input().split()))
    C = list(map(int, input().split()))
    
    # Count how many times each pot is exposed to sunlight
    # and track the order of operations for each pot
    pot_operations = defaultdict(list)
    for j, c in enumerate(C):
        pot_operations[c - 1].append(j)  # 0-indexed
    
    # For each pot, we need to track:
    # - When it receives sunlight, water evaporates and flows down
    # - When water flows in from above
    
    # Let's simulate but optimize by processing in order
    # Key insight: water only flows downward, so we can process from top to bottom
    
    # For each pot i, we need to know all times it's exposed to sunlight
    # Between consecutive exposures, water can accumulate from pot i-1
    
    # Let's think differently:
    # Process all Q operations, but track water flow efficiently
    
    # Actually, let's just simulate but be smart about it
    # When pot i is exposed, water flows to pot i+1
    # We need to handle this correctly
    
    # Simple simulation should work for N, Q <= 2*10^5 if we're careful
    # Each operation only affects one pot directly, but water flows down
    
    water = A[:]
    
    for c in C:
        i = c - 1  # 0-indexed
        current_water = water[i]
        evaporated = min(current_water, B[i])
        remaining = current_water - evaporated
        water[i] = 0
        if i + 1 < N:
            water[i + 1] += remaining
        # if i + 1 == N, water goes to drain (disappears)
    
    print(' '.join(map(str, water)))

if __name__ == '__main__':
    main()

この解説は claude4.5opus によって生成されました。

posted:
last update: