题意:给定一个字符串,字符串只包括
+
,
−
+,-
+,−,如果一个字符串中的
+
+
+的数量和
−
-
−的数量相等,我们就称为是平衡的字符串,如果能通过以下操作使得字符串变成平衡,我们就说该字符串是有希望的,平衡的字符串一定有希望。问一个字符串有多少子串是有希望的?
操作:可以用相邻的两个 − - −替换成 + + +
思路:记一个子串中的
+
+
+的个数为b,
−
-
−的个数为a,可以由
−
-
−转换成
+
+
+的个数为k,那么就有
a
−
2
∗
k
=
b
+
k
a-2*k=b+k
a−2∗k=b+k
即
a
−
b
=
3
k
a-b=3k
a−b=3k
现在我们令
+
+
+为
−
1
-1
−1,令
−
-
−为
+
1
+1
+1,用
p
r
e
[
i
]
pre[i]
pre[i]表示前缀和,因为是由
−
-
−转换成
+
+
+,那么题目就转换成了
p
r
e
[
i
]
−
p
r
e
[
j
]
≥
0
pre[i]-pre[j] \ge 0
pre[i]−pre[j]≥0 并且有
(
p
r
e
[
i
]
−
p
r
e
[
j
]
)
%
3
=
=
0
(
i
>
j
)
(pre[i]-pre[j])\%3 == 0 (i>j)
(pre[i]−pre[j])%3==0(i>j)
为什么要令 + + +为 − 1 -1 −1,令 − - −为 + 1 +1 +1呢,这样做其实就保证了 − - −是连续出现的,如果不是连续出现的,其前缀和是不会达到3的。这一点可以自己举个例证明一下
为了满足这两个不等式,我们可以用三个树状数组分别维护
p
r
e
[
i
]
%
3
pre[i]\%3
pre[i]%3 的值,然后用树状数组求比
p
r
e
[
i
]
pre[i]
pre[i] 小的个数。这里有个小细节! 因为用树状数组求比一个数小的个数不能出现负数,前缀和中可能会出现负数,因此我们首先计算出前缀和的最小值
m
n
mn
mn,然后让每个数都增加
m
n
+
1
mn+1
mn+1。
还有一个小细节就是
p
r
e
[
0
]
pre[0]
pre[0]也要加入到树状数组中去。如果不把0加入到树状数组中去的话,那么会漏掉第一次出现3的情况
代码如下:
#include <iostream>
#include <algorithm>
#include <cstring>
#include <string>
#include <vector>
#include <map>
#include <cmath>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
const int N = 5e5 + 10;//注意大小
int n;
int pre[N];
int tr[3][N];
int lowbit(int x)
{
return x&-x;
}
void add(int x,int c,int mod)
{
for(int i=x;i<N;i+=lowbit(i))tr[mod][i]+=c;//注意边界
}
int sum(int x,int mod)
{
int cnt=0;
for(int i=x;i>=1;i-=lowbit(i))cnt+=tr[mod][i];
return cnt;
}
int main()
{
ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
int T;cin>>T;
while(T--)
{
cin>>n;
string s;cin>>s;s=" "+s;
int mn=0;//前缀和中的最小值
pre[0]=0;//注意0位置要重定义一下
for(int i=1;i<=n;i++)
{
pre[i]=pre[i-1]+(s[i]=='+'?-1:1);
mn=min(mn,pre[i]);
}
for(int i=0;i<=n;i++)pre[i]+=abs(mn)+1;//将所有的前缀和都增加mn
for(int i=0;i<3;i++)//初始化
for(int j=0;j<=abs(mn)+n+1;j++)//注意边界
tr[i][j]=0;
ll ans=0;//树状数组求一个数比他小的个数
for(int i=0;i<=n;i++)
{
int mod=pre[i]%3;//维护三个pre[i]%3的数组
ans+=sum(pre[i],mod);
add(pre[i],1,mod);
}
cout<<ans<<"\n";
}
return 0;
}