Official

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

Gemini 3.0 Flash (Thinking)

概要

\(N\) 日間の食事記録から、野菜スコアの合計と果物スコアの合計が等しくなる区間(バランス期間)の総数を求める問題です。

考察

条件の言い換え

ある区間 \([l, r]\) において、「野菜スコアの合計 \(\sum v_i\)」と「果物スコアの合計 \(\sum f_i\)」が等しいという条件は、次のように言い換えることができます。 $\(\sum_{i=l}^{r} v_i = \sum_{i=l}^{r} f_i \iff \sum_{i=l}^{r} (v_i - f_i) = 0\)$

ここで、各日の「野菜スコアと果物スコアの差」を \(d_i = v_i - f_i\) と定義すると、各文字に対応する \(d_i\) の値は以下のようになります。 - V のとき:\(d_i = 1 - 0 = 1\) - F のとき:\(d_i = 0 - 1 = -1\) - B のとき:\(d_i = 1 - 1 = 0\) - N のとき:\(d_i = 0 - 0 = 0\)

したがって、問題は「連続する部分列の和が \(0\) になる区間の個数を数える」という標準的な問題に帰着されます。

累積和の利用

区間 \([l, r]\) の和が \(0\) であることを判定するために、累積和 \(D_k = \sum_{i=1}^k d_i\)(ただし \(D_0 = 0\))を考えます。 区間 \([l, r]\) の和は \(D_r - D_{l-1}\) と表せるため、条件は以下のようになります。 $\(D_r - D_{l-1} = 0 \iff D_r = D_{l-1}\)$

つまり、累積和の配列 \(D_0, D_1, \dots, D_N\) の中で、同じ値を持つインデックスのペア \((i, j) \ (i < j)\) の個数を数えればよいことになります。

アルゴリズム

  1. 累積和を計算するための変数 cur\(0\) で初期化します。
  2. 各累積和の値が何回出現したかを記録する連想配列(または配列)を用意します。
  3. 最初に \(D_0 = 0\) に対応する出現回数を \(1\) 増やしておきます。
  4. 文字列 \(S\) を先頭から順に見ていき、各文字に応じて cur を更新(Vなら \(+1\)Fなら \(-1\)、それ以外は変化なし)します。
  5. 更新後の cur の値を記録していきます。
  6. 各値の出現回数を \(c\) とすると、その値から作れるペアの数は \(\displaystyle \frac{c(c-1)}{2}\) となります。これをすべての値について合計したものが答えです。

計算量

  • 時間計算量: \(O(N)\)
    • 文字列の走査に \(O(N)\)、出現回数の集計と答えの計算に \(O(N)\) かかります。
  • 空間計算量: \(O(N)\)
    • 出現回数を記録する配列(または連想配列)に最大で \(2N+1\) 個の要素を格納するため、\(O(N)\) となります。

実装のポイント

  • 負のインデックスの扱い: 累積和の値は \(-N\) から \(N\) の範囲を取り得ます。Python の dict を使う場合は問題ありませんが、配列で管理する場合は、値に \(N\) を足してオフセットを加えることで、インデックスを \(0\) 以上の範囲(\(0\) から \(2N\))に収める必要があります。

  • 答えの大きさ: 答えは最大で \(\frac{N(N+1)}{2}\) 程度(\(N=10^6\) のとき約 \(5 \times 10^{11}\))になるため、32bit整数型ではオーバーフローする可能性があります。Python では自動的に多倍長整数として扱われるため気にする必要はありませんが、他の言語(C++など)では long long 型などを使用してください。

    ソースコード

import sys

def solve():
    # 標準入力から全てのデータを読み込み、空白で分割します。
    # N は日数、S は食事記録の文字列です。
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    n = int(input_data[0])
    s = input_data[1]
    
    # 野菜スコアを v_i、果物スコアを f_i としたとき、
    # 区間 [l, r] がバランス期間である条件は sum_{i=l}^r v_i = sum_{i=l}^r f_i です。
    # これは sum_{i=l}^r (v_i - f_i) = 0 と言い換えられます。
    # 各文字における (v_i - f_i) の値は以下の通りです:
    # 'V' (野菜のみ): 1 - 0 = 1
    # 'F' (果物のみ): 0 - 1 = -1
    # 'B' (両方): 1 - 1 = 0
    # 'N' (どちらも摂らない): 0 - 0 = 0
    
    # 累積和 D_k = sum_{i=1}^k (v_i - f_i) を定義すると、
    # 条件は D_r - D_{l-1} = 0、すなわち D_r = D_{l-1} となります。
    
    # 累積和 D_k の出現回数をカウントするための配列を用意します。
    # D_k の範囲は -N から N までなので、サイズ 2N + 1 の配列を用い、
    # インデックスにオフセット N を加えることで負の値を扱います。
    counts = [0] * (2 * n + 1)
    
    # 初期状態 D_0 = 0 (オフセット込みでインデックス n)
    cur = n
    counts[cur] = 1
    
    # 文字列を 1 文字ずつ走査し、累積和を更新しながらカウントします。
    for char in s:
        if char == 'V':
            cur += 1
        elif char == 'F':
            cur -= 1
        # 'B' と 'N' の場合は cur の値は変わりません。
        
        counts[cur] += 1
        
    # 各累積和の値について、出現回数を c とすると、
    # その値を持つインデックスのペア (i, j) [i < j] の個数は cC2 = c * (c - 1) / 2 です。
    ans = 0
    for c in counts:
        if c > 1:
            ans += c * (c - 1) // 2
            
    # 結果を出力します。
    sys.stdout.write(str(ans) + '\n')

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-thinking によって生成されました。

posted:
last update: