公式

D - バランス食生活 / Balanced Diet 解説 by admin

Claude 4.6 Opus (Thinking)

概要

\(N\) 日間の食事記録から、野菜スコアの合計と果物スコアの合計が等しくなる連続区間 \((l, r)\) の個数を求める問題です。累積和の差分に注目することで効率的に解くことができます。

考察

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

すべての組 \((l, r)\) を列挙し、それぞれの区間について野菜スコアと果物スコアの合計を計算すると、組の数が \(O(N^2)\) 個あるため、\(N \leq 10^6\) では到底間に合いません。

重要な気づき:差分に注目する

バランス期間の条件は次の通りです:

\[\sum_{i=l}^{r} v_i = \sum_{i=l}^{r} f_i\]

これを変形すると:

\[\sum_{i=l}^{r} (v_i - f_i) = 0\]

ここで、各文字に対する \(d_i = v_i - f_i\) の値を考えます:

文字 \(v_i\) \(f_i\) \(d_i = v_i - f_i\)
V \(1\) \(0\) \(+1\)
F \(0\) \(1\) \(-1\)
B \(1\) \(1\) \(0\)
N \(0\) \(0\) \(0\)

累積和への帰着

\(d_i\) の累積和(プレフィックス和)を \(P[k] = \sum_{i=1}^{k} d_i\)\(P[0] = 0\))と定義すると、区間 \([l, r]\) の条件は:

\[P[r] - P[l-1] = 0 \quad \Longleftrightarrow \quad P[r] = P[l-1]\]

つまり、累積和の値が同じになる2つのインデックスの組を数える問題 に帰着できます。

具体例

\(S = \) VFBV のとき、\(d = [+1, -1, 0, +1]\) で、累積和は \(P = [0, 1, 0, 0, 1]\) となります。

  • \(P[0] = P[2] = P[3] = 0\) → この3つから \(\binom{3}{2} = 3\)
  • \(P[1] = P[4] = 1\) → この2つから \(\binom{2}{2} = 1\)

合計 \(4\) 個のバランス期間があります。

アルゴリズム

  1. 累積和 \(P[0] = 0\) から開始する。
  2. 文字列 \(S\) を1文字ずつ走査し、V なら \(+1\)F なら \(-1\)BN なら \(0\) を累積和に加算する。
  3. 各ステップで、現在の累積和の値がこれまでに何回出現したかをハッシュマップ(辞書)で管理する。
  4. 現在の累積和と同じ値が過去に \(c\) 回出現していれば、\(c\) 個のバランス期間が新たに見つかるので、答えに \(c\) を加算する。
  5. 最終的な答えを出力する。

このテクニックは「同じ累積和の値のペアを数える」典型的な手法で、逐次的にカウントを更新することで \(O(N)\) で処理できます。

計算量

  • 時間計算量: \(O(N)\) — 文字列を1回走査するだけで済みます。
  • 空間計算量: \(O(N)\) — 累積和の値は \(-N\) から \(N\) の範囲なので、最悪 \(O(N)\) 個のエントリが辞書に格納されます。

実装のポイント

  • 累積和の初期値 \(P[0] = 0\) を忘れずにカウントに含める(count[0] = 1 で初期化)。これを忘れると、\(l = 1\) から始まるバランス期間を見逃してしまいます。

  • BN は差分が \(0\) なので累積和を変化させないことに注意。特別な処理は不要です。

  • \(N\) が最大 \(10^6\) と大きいため、Python では sys.stdin.buffer.read() で高速に入力を読み取っています。

    ソースコード

import sys
from collections import defaultdict

def main():
    input_data = sys.stdin.buffer.read().decode()
    tokens = input_data.split()
    N = int(tokens[0])
    S = tokens[1]
    
    # v_i - f_i の累積和を考える
    # V: v=1, f=0 -> diff = +1
    # F: v=0, f=1 -> diff = -1
    # B: v=1, f=1 -> diff = 0
    # N: v=0, f=0 -> diff = 0
    # バランス期間: sum(v_i, l..r) == sum(f_i, l..r)
    # <=> prefix_diff[r] - prefix_diff[l-1] == 0
    # <=> prefix_diff[r] == prefix_diff[l-1]
    
    # Count pairs with same prefix_diff value
    count = defaultdict(int)
    prefix = 0
    count[0] = 1
    result = 0
    
    for ch in S:
        if ch == 'V':
            prefix += 1
        elif ch == 'F':
            prefix -= 1
        # B and N: prefix unchanged
        
        result += count[prefix]
        count[prefix] += 1
    
    print(result)

main()

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

投稿日時:
最終更新: