公式

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

GPT 5.4 High

概要

各日について「野菜スコア \(-\) 果物スコア」の差だけを見ると、区間がバランス期間である条件は「その区間の差の総和が \(0\)」と言い換えられます。
そこで、差の累積和(prefix sum)が同じになる位置の組を数えることで、全てのバランス期間を \(O(N)\) で数えられます。

考察

まず、各文字がバランスにどう影響するかを考えます。

  • V : 野菜だけなので差は \(+1\)
  • F : 果物だけなので差は \(-1\)
  • B : 両方なので差は \(0\)
  • N : どちらもないので差は \(0\)

つまり、各日 \(i\) について

\[ d_i = v_i - f_i \]

とおくと、

文字 \(d_i\)
V \(1\)
F \(-1\)
B \(0\)
N \(0\)

となります。

区間 \([l, r]\) がバランス期間である条件は

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

です。これを移項すると

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

つまり

\[ \sum_{i=l}^{r} d_i = 0 \]

となります。


ここで累積和を使います。
先頭から \(i\) 日目までの差の累積和を

\[ D_i = \sum_{j=1}^{i} d_j \]

とし、さらに \(D_0 = 0\) とします。

すると区間 \([l, r]\) の差の総和は

\[ \sum_{i=l}^{r} d_i = D_r - D_{l-1} \]

なので、区間 \([l, r]\) がバランス期間であることは

\[ D_r = D_{l-1} \]

と同値です。

つまりこの問題は、

累積和 \(D_0, D_1, \dots, D_N\) の中で、同じ値になる 2 つの位置の組の個数を数える問題

に変わります。


素朴な方法がだめな理由

全ての区間 \((l, r)\) を試すと、区間数は約 \(N^2/2\) 個あります。
\(N \le 10^6\) なので、\(O(N^2)\) では到底間に合いません。


どう高速化するか

左から順に累積和を見ていくとき、現在の累積和と同じ値が以前に何回現れたかが分かれば、その回数だけ新しいバランス期間ができます。

たとえば現在の累積和が \(x\) で、過去に \(x\)\(k\) 回出ていたなら、その \(k\) 個それぞれに対して区間を作れるので、答えに \(k\) を足せばよいです。

この「各累積和が何回現れたか」を配列や連想配列で管理すれば、1 文字見るごとに \(O(1)\) で更新できます。


具体例

例えば \(S = \texttt{VBNF}\) とします。

各日の差 \(d_i\)

  • V \(\to 1\)
  • B \(\to 0\)
  • N \(\to 0\)
  • F \(\to -1\)

なので、累積和は

\[ D_0 = 0,\quad D_1 = 1,\quad D_2 = 1,\quad D_3 = 1,\quad D_4 = 0 \]

です。

同じ値が出ている場所を見ると、

  • \(0\)\(D_0, D_4\) の 2 回
  • \(1\)\(D_1, D_2, D_3\) の 3 回

同じ値が \(m\) 回出たら、そこから作れる組の数は

\[ \binom{m}{2} \]

です。したがって

  • \(0\) から \(1\)
  • \(1\) から \(3\)

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

アルゴリズム

  1. 差の累積和 diff を管理する。初期値は \(0\)
  2. cnt[x] を「累積和 \(x\) がこれまでに何回現れたか」とする。
  3. 最初に \(D_0 = 0\) が 1 回あるので、cnt[0] = 1 としておく。
  4. 文字列を左から 1 文字ずつ見る。
    • V なら diff += 1
    • F なら diff -= 1
    • B, N なら変化なし
  5. 現在の diff が過去に cnt[diff] 回出ていれば、その回数だけ新しいバランス期間があるので、答えに足す。
  6. その後、cnt[diff] を 1 増やす。
  7. 最後に答えを出力する。

このコードで配列を使っている理由

累積和 diff は 1 日ごとに \(+1, 0, -1\) しか変化しないので、値の範囲は必ず

\[ -N \le diff \le N \]

に収まります。

そのため、辞書ではなく長さ \(2N+1\) の配列を使って頻度を管理できます。
負の添字を避けるために offset = N を足して、

\[ \text{index} = diff + N \]

として配列にアクセスしています。

計算量

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

実装のポイント

  • cnt[offset] = 1 として、最初の累積和 \(0\) を忘れずに入れることが重要です。これがないと、先頭から始まるバランス期間を数え漏らします。

  • BN はどちらも差が \(0\) なので、diff を変化させません。

  • 答えの個数は最大で \(O(N^2)\) になり得ますが、Python の整数ならそのまま扱えます。

  • 提示コードでは入力を bytes として読み込んでいるため、文字比較は ord('V') のように整数で行っています。

    ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    N = int(data[0])
    S = data[1]

    offset = N
    cnt = [0] * (2 * N + 1)
    cnt[offset] = 1

    diff = 0
    ans = 0

    for c in S:
        if c == ord('V'):
            diff += 1
        elif c == ord('F'):
            diff -= 1
        idx = offset + diff
        ans += cnt[idx]
        cnt[idx] += 1

    print(ans)

if __name__ == "__main__":
    main()

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

投稿日時:
最終更新: