Official

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


以下、数列の添字は \(0\) から始まるものとします。

野菜スコアと果物スコアが等しいことは、「野菜スコア-果物スコア=0」と同値です。よってそのような区間を探すことを考えます。

「野菜スコア-果物スコア」は B, N では変わらず、’V’ では1増え、F では 1減ります。よって問題は次のように読み替えられます:

\(S\)B, N\(0\) に、V\(1\) に、F\(-1\) に置き換えてできる数列を \(A\) とします。\(A\) の連続部分列のうち、和が \(0\) になるものの個数は?

ここで、数列 \(B\)\(A\) の累積和とします。すなわち

\(B_i=\begin{cases} 0 & i=0 のとき\\ B_{i-1}+A_{i-1} & i>0 のとき \end{cases}\)

と定めます。\(B_i\) は、数列 \(A\) の先頭から \(i\) 項の和になります。定義から \(B\)\(O(N)\) で求めることができます。

このとき、\(A\) の連続する \(M\) 要素の和 \(A_L+A_{L+1}+\dots+A_{R-1}\)\(B_{R}-B_{L}\) となるので \(A_L+\dots+A_R=0\)\(B_R-B_L=0\) と同値です。

よって問題は更に次のように読み替えられます:

\(0 \leq i<j \leq N\) を満たす整数の組 \((i,j)\) のうち、 \(B_i=B_j\) を満たすものは何個か?

この問題は \(S\) を前から順に確かめながら、どの値が何度登場したかを辞書で記録することで、 \(O(N\log N)\) で解くことができます。今回は値の範囲が小さいため、辞書ではなく単に配列で保持してもよく、\(O(N)\) で解くこともできます。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n;
  cin >> n;
  string s;
  cin >> s;

  map<int,int>d;
  long long ans = 0;
  int bi = 0;
  d[bi]++;
  for(int i=0; i<n; i++){
    if(s[i] == 'F'){
      bi++;
    }else if(s[i] == 'V'){
      bi--;
    }
    ans += d[bi];
    d[bi]++;
  }

  cout << ans << endl;
}

実装例 (Python)

from collections import defaultdict

N = int(input())
S = input()

d = defaultdict(int)
ans = 0
bi = 0
d[bi] += 1
for c in S:
  if c == 'F':
    bi += 1
  elif c == 'V':
    bi -= 1
  ans += d[bi]
  d[bi] += 1

print(ans)

posted:
last update: