公式

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

Gemini 3.0 Flash (Thinking)

Overview

Given a meal record over \(N\) days, the problem asks to find the total number of intervals (balanced periods) where the sum of vegetable scores equals the sum of fruit scores.

Analysis

Rephrasing the Condition

For an interval \([l, r]\), the condition that “the sum of vegetable scores \(\sum v_i\)” equals “the sum of fruit scores \(\sum f_i\)” can be rephrased as follows: $\(\sum_{i=l}^{r} v_i = \sum_{i=l}^{r} f_i \iff \sum_{i=l}^{r} (v_i - f_i) = 0\)$

Here, if we define the “difference between vegetable score and fruit score” for each day as \(d_i = v_i - f_i\), the value of \(d_i\) corresponding to each character is as follows: - V: \(d_i = 1 - 0 = 1\) - F: \(d_i = 0 - 1 = -1\) - B: \(d_i = 1 - 1 = 0\) - N: \(d_i = 0 - 0 = 0\)

Therefore, the problem reduces to the standard problem of “counting the number of intervals where the sum of a contiguous subarray is \(0\).”

Using Prefix Sums

To determine whether the sum of an interval \([l, r]\) is \(0\), we consider the prefix sum \(D_k = \sum_{i=1}^k d_i\) (where \(D_0 = 0\)). Since the sum of the interval \([l, r]\) can be expressed as \(D_r - D_{l-1}\), the condition becomes: $\(D_r - D_{l-1} = 0 \iff D_r = D_{l-1}\)$

In other words, we need to count the number of pairs of indices \((i, j) \ (i < j)\) that have the same value in the prefix sum array \(D_0, D_1, \dots, D_N\).

Algorithm

  1. Initialize a variable cur to \(0\) for computing the prefix sum.
  2. Prepare an associative array (or array) to record how many times each prefix sum value has appeared.
  3. First, increment the count for \(D_0 = 0\) by \(1\).
  4. Iterate through the string \(S\) from the beginning, updating cur according to each character (V adds \(+1\), F adds \(-1\), others cause no change).
  5. Record the value of cur after each update.
  6. If the number of occurrences of a particular value is \(c\), the number of pairs that can be formed from that value is \(\displaystyle \frac{c(c-1)}{2}\). The answer is the sum of this over all values.

Complexity

  • Time Complexity: \(O(N)\)
    • Scanning the string takes \(O(N)\), and tallying the occurrence counts and computing the answer takes \(O(N)\).
  • Space Complexity: \(O(N)\)
    • The array (or associative array) for recording occurrence counts stores at most \(2N+1\) elements, resulting in \(O(N)\).

Implementation Notes

  • Handling negative indices: The prefix sum values can range from \(-N\) to \(N\). If using Python’s dict, this is not an issue, but if managing with an array, you need to add an offset of \(N\) to the values so that the indices fit within the non-negative range (\(0\) to \(2N\)).

  • Size of the answer: The answer can be as large as approximately \(\frac{N(N+1)}{2}\) (about \(5 \times 10^{11}\) when \(N=10^6\)), so it may overflow a 32-bit integer type. In Python, this is not a concern since integers are automatically handled as arbitrary-precision, but in other languages (such as C++), please use long long or similar types.

    Source Code

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

This editorial was generated by gemini-3-flash-thinking.

投稿日時:
最終更新: