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\) 個のバランス期間があります。
アルゴリズム
- 差の累積和
diffを管理する。初期値は \(0\)。 cnt[x]を「累積和 \(x\) がこれまでに何回現れたか」とする。- 最初に \(D_0 = 0\) が 1 回あるので、
cnt[0] = 1としておく。 - 文字列を左から 1 文字ずつ見る。
Vならdiff += 1Fならdiff -= 1B,Nなら変化なし
- 現在の
diffが過去にcnt[diff]回出ていれば、その回数だけ新しいバランス期間があるので、答えに足す。 - その後、
cnt[diff]を 1 増やす。 - 最後に答えを出力する。
このコードで配列を使っている理由
累積和 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\) を忘れずに入れることが重要です。これがないと、先頭から始まるバランス期間を数え漏らします。BとNはどちらも差が \(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 によって生成されました。
投稿日時:
最終更新: