Official

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

Claude 4.6 Opus (Thinking)

Overview

Given a meal record over \(N\) days, the problem asks to find the number of contiguous intervals \((l, r)\) where the total vegetable score equals the total fruit score. This can be solved efficiently by focusing on the difference of prefix sums.

Analysis

Naive Approach and Its Issues

If we enumerate all pairs \((l, r)\) and compute the vegetable and fruit score totals for each interval, the number of pairs is \(O(N^2)\), which is far too slow for \(N \leq 10^6\).

Key Insight: Focus on the Difference

The condition for a balanced period is:

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

Rearranging this:

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

Here, consider the value \(d_i = v_i - f_i\) for each character:

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

Reduction to Prefix Sums

Define the prefix sum of \(d_i\) as \(P[k] = \sum_{i=1}^{k} d_i\) (with \(P[0] = 0\)). Then the condition for interval \([l, r]\) becomes:

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

In other words, the problem reduces to counting pairs of indices that have the same prefix sum value.

Concrete Example

For \(S = \) VFBV, we have \(d = [+1, -1, 0, +1]\), and the prefix sums are \(P = [0, 1, 0, 0, 1]\).

  • \(P[0] = P[2] = P[3] = 0\) → From these 3 indices, \(\binom{3}{2} = 3\) pairs
  • \(P[1] = P[4] = 1\) → From these 2 indices, \(\binom{2}{2} = 1\) pair

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

Algorithm

  1. Start with prefix sum \(P[0] = 0\).
  2. Scan the string \(S\) character by character, adding \(+1\) for V, \(-1\) for F, and \(0\) for B or N to the prefix sum.
  3. At each step, use a hash map (dictionary) to track how many times each prefix sum value has appeared so far.
  4. If the current prefix sum value has appeared \(c\) times in the past, then \(c\) new balanced periods are found, so add \(c\) to the answer.
  5. Output the final answer.

This technique is a classic method for “counting pairs with the same prefix sum value,” and by incrementally updating the counts, it runs in \(O(N)\).

Complexity

  • Time complexity: \(O(N)\) — Only a single pass through the string is needed.
  • Space complexity: \(O(N)\) — The prefix sum values range from \(-N\) to \(N\), so at most \(O(N)\) entries are stored in the dictionary.

Implementation Notes

  • Don’t forget to include the initial prefix sum value \(P[0] = 0\) in the count (initialize with count[0] = 1). Forgetting this will cause balanced periods starting from \(l = 1\) to be missed.

  • Note that B and N have a difference of \(0\), so they do not change the prefix sum. No special handling is needed.

  • Since \(N\) can be as large as \(10^6\), in Python we use sys.stdin.buffer.read() for fast input reading.

    Source Code

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

This editorial was generated by claude4.6opus-thinking.

posted:
last update: