Official

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

GPT 5.4 High

Overview

If we look only at the difference “vegetable score \(-\) fruit score” for each day, the condition for an interval to be a balanced period can be rephrased as “the sum of differences over that interval is \(0\).”
By counting pairs of positions where the prefix sum of differences takes the same value, we can count all balanced periods in \(O(N)\).

Analysis

First, let’s consider how each character affects the balance.

  • V : vegetables only, so the difference is \(+1\)
  • F : fruits only, so the difference is \(-1\)
  • B : both, so the difference is \(0\)
  • N : neither, so the difference is \(0\)

In other words, for each day \(i\), if we define

\[ d_i = v_i - f_i \]

then:

Character \(d_i\)
V \(1\)
F \(-1\)
B \(0\)
N \(0\)

The condition for an interval \([l, r]\) to be a balanced period is

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

Rearranging this gives

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

That is,

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


Here we use prefix sums.
Let the prefix sum of differences up to day \(i\) be

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

and additionally define \(D_0 = 0\).

Then the sum of differences over interval \([l, r]\) is

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

so the interval \([l, r]\) being a balanced period is equivalent to

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

In other words, this problem transforms into:

Counting the number of pairs of positions among the prefix sums \(D_0, D_1, \dots, D_N\) that have the same value.


Why the naive approach doesn’t work

If we try all intervals \((l, r)\), there are approximately \(N^2/2\) intervals.
Since \(N \le 10^6\), an \(O(N^2)\) approach is far too slow.


How to speed it up

As we scan the prefix sums from left to right, if we know how many times the current prefix sum value has appeared before, that count equals the number of new balanced periods we can form.

For example, if the current prefix sum is \(x\) and \(x\) has appeared \(k\) times in the past, we can form an interval with each of those \(k\) occurrences, so we add \(k\) to the answer.

By maintaining “how many times each prefix sum value has appeared” using an array or associative array, we can update in \(O(1)\) per character.


Concrete example

For example, let \(S = \texttt{VBNF}\).

The difference \(d_i\) for each day is:

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

So the prefix sums are

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

Looking at positions with the same value:

  • Value \(0\) appears at \(D_0, D_4\) — 2 times
  • Value \(1\) appears at \(D_1, D_2, D_3\) — 3 times

If the same value appears \(m\) times, the number of pairs that can be formed is

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

Therefore:

  • From value \(0\): \(1\) pair
  • From value \(1\): \(3\) pairs

In total, there are \(4\) balanced periods.

Algorithm

  1. Maintain the prefix sum of differences diff. Initialize it to \(0\).
  2. Let cnt[x] be “how many times prefix sum \(x\) has appeared so far.”
  3. Since \(D_0 = 0\) occurs once at the start, set cnt[0] = 1.
  4. Scan the string from left to right, one character at a time.
    • If V, then diff += 1
    • If F, then diff -= 1
    • If B or N, no change
  5. If the current diff has appeared cnt[diff] times before, that many new balanced periods exist, so add that count to the answer.
  6. Then increment cnt[diff] by 1.
  7. Finally, output the answer.

Why the code uses an array

Since the prefix sum diff changes by only \(+1, 0, -1\) per day, its range is always

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

Therefore, instead of a dictionary, we can use an array of length \(2N+1\) to manage frequencies.
To avoid negative indices, we add offset = N, accessing the array with

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

Complexity

  • Time complexity: \(O(N)\)
  • Space complexity: \(O(N)\)

Implementation notes

  • It is important to set cnt[offset] = 1 to account for the initial prefix sum of \(0\). Without this, balanced periods starting from the beginning would be missed.

  • Both B and N have a difference of \(0\), so they do not change diff.

  • The answer can be as large as \(O(N^2)\), but Python’s integers can handle this directly.

  • In the provided code, the input is read as bytes, so character comparisons are done using integers like ord('V').

    Source Code

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()

This editorial was generated by gpt-5.4-high.

posted:
last update: