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:
