Official

B - 芝生の水やり / Watering the Lawn Editorial by admin

Claude 4.5 Opus

概要

連続する \(K\) 区画に同時に水をまく操作を繰り返して、各区画の水分量をちょうど目標値に一致させられるかを判定する問題です。

考察

重要な気づき

この問題のポイントは、左から順番に見ていくと、各位置での「水をまく回数」が一意に決まるということです。

具体例で考えてみましょう。\(N=5, K=3\)、目標値が \([2, 3, 4, 2, 1]\) の場合:

  • 区画0の水分量を \(2\) にするには、位置0を左端として2回水をまく必要があります
  • 区画1には、位置0からの水やり(2回)が影響しています。目標は3なので、位置1から追加で1回水をまきます
  • このように左から順に、必要な回数を決定していきます

素朴なアプローチの問題点

各区画への水分量を計算するたびに、影響を与える全ての水やり回数を足し合わせると、最悪 \(O(NK)\) の計算量になりTLEします。

解決策

スライディングウィンドウを使います。区画 \(i\) に影響を与える水やりは、位置 \(\max(0, i-K+1)\) から位置 \(i\) までの \(K\) 個です。これは連続する区間なので、区画を1つ進むたびに: - 新しく追加される水やり回数を足す - 区間から外れる水やり回数を引く

これにより、各区画の水分量を \(O(1)\) で計算できます。

不可能になる条件

  1. すでに目標を超えている: 現在の水分量が目標値より大きい場合(水を引くことはできない)
  2. 調整できない: 位置 \(i\) を左端として水をまけるのは \(i + K - 1 < N\)(つまり \(i \leq N - K\))のときのみ。それ以降は新たに水やりを追加できないので、水分量が目標と一致していなければ不可能

アルゴリズム

  1. 配列 \(water[i]\) = 位置 \(i\) を左端として水をまく回数、を管理
  2. \(current\_sum\) = 現在の区画に影響している水やり回数の合計(スライディングウィンドウ)
  3. 左から順に各区画 \(i\) を処理:
    • \(needed = T[i] - current\_sum\)(追加で必要な水分量)
    • \(needed < 0\) なら No(すでに超過)
    • \(i \leq N - K\) なら \(water[i] = needed\) として \(current\_sum\) に加算
    • \(i > N - K\)\(needed \neq 0\) なら No(調整不能)
    • ウィンドウから外れる \(water[i-K+1]\)\(current\_sum\) から引く
  4. 全て処理できたら Yes

計算量

  • 時間計算量: \(O(N)\) — 各区画を1回ずつ処理
  • 空間計算量: \(O(N)\) — 配列 \(water\) と入力配列 \(T\) の保存

実装のポイント

  • スライディングウィンドウの更新タイミングに注意。区画 \(i+1\) に移る前に、位置 \(i-K+1\) の影響を除去します

  • 位置 \(i\) を左端として水をまけるのは \(i \leq N - K\) のときのみ。境界条件を正確に扱うことが重要

  • \(needed\) が負になるケースは、前の水やりで水分量が超過してしまった場合なので、即座に No を返します

    ソースコード

def solve():
    N, K = map(int, input().split())
    T = list(map(int, input().split()))
    
    # 各位置での水やり回数を決定していく
    # water[i] = 位置iを左端として水をまく回数
    # 位置iの水分量 = sum(water[max(0, i-K+1):i+1])
    
    # imos法的に、現在の水分量を追跡
    # current_water = 現在の区画に与えられている水分量
    
    # water[i]を決めるとき、区画iの水分量は
    # water[i-K+1] + water[i-K+2] + ... + water[i] で決まる
    # 区画0から順に見ていくと、区画iを見るときにwater[i]を決定できる
    
    water = [0] * N  # water[i] = 位置iを左端として水をまく回数
    current_sum = 0  # 現在の区画に影響を与えている水やりの合計
    
    for i in range(N):
        # 区画iに対する水分量は、位置max(0, i-K+1)からiまでの水やり回数の合計
        # current_sumにはその合計が入っている(スライディングウィンドウ)
        
        # 区画iの目標水分量はT[i]
        # 現在の水分量はcurrent_sum
        # 追加で必要な水分量はT[i] - current_sum
        # これをwater[i]で賄う(位置iを左端として水をまく)
        
        needed = T[i] - current_sum
        
        if needed < 0:
            # すでに目標を超えている
            print("No")
            return
        
        # 位置iを左端として水をまく回数を設定
        # ただし、位置iを左端にできるのは i + K - 1 < N、つまり i <= N - K のときのみ
        if i <= N - K:
            water[i] = needed
            current_sum += needed
        else:
            # 位置iを左端として水をまくことはできない
            # なので、current_sumがちょうどT[i]でなければならない
            if needed != 0:
                print("No")
                return
        
        # 次の区画に移る前に、ウィンドウから外れる分を引く
        # 区画i+1に影響を与えるのは位置max(0, i+1-K+1)からi+1まで
        # 区画iに影響を与えていたのは位置max(0, i-K+1)からiまで
        # 位置i-K+1がウィンドウから外れる
        if i - K + 1 >= 0:
            current_sum -= water[i - K + 1]
    
    print("Yes")

solve()

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

posted:
last update: