前缀和
十分邪恶的区间子区间问题,还是先理解题意。
给定一个 $\texttt{01}$ 串 $s$,定义 $f(l,r)$ 为子串 $[l,r]$ 中 $\texttt{0}$ 与 $\texttt{1}$ 的数量相等的子串数量,求下式。
$$\left( \sum\limits_{l=1}^{n} \sum\limits_{r=l}^{n} f(l,r) \right) \bmod (10^9 + 7)$$
简单转化一下,把 $0$ 改成 $-1$,那么子串中 $\texttt{01}$ 数量相等当且仅当区间和为 $0$。考虑从 $1$ 到 $n$ 枚举 $i$ 并统一计算以 $i$ 为右端点的所有区间的答案。
继续转化,使用 $sum_i$ 表示位置 $i$ 处的前缀和,试图找到一种好的表示答案的方法。这里考虑对于每个和为 $0$ 的区间计算贡献。
$$\sum\limits_{l=1}^{n} \sum\limits_{r=l}^{n} l \times (n-r+1) \times [sum_{l-1} = sum_{r}]$$
得到启发,现在只考虑前缀和相等的位置。对于右端点 $r$ 记 $w_{sum_r}$ 为 $sum_i = sum_r$ 且 $i < r$ 的整数 $i$ 的 $i+1$ 的和,注意计入 $i=0$ 的情况。那么以 $r$ 作为右端点的 $\texttt{01}$ 数量相等的区间对答案的贡献就是 $(n-r+1) \times w_{sum_r}$,每次计算之后更新 $w_{sum_r}$ 即可。
使用 map
维护序列 $w$,时间复杂度 $O(n \log n)$。
#include<bits/stdc++.h>
using namespace std;
constexpr int mod=1'000'000'007;
int n,ans,sum[200'005];
void solution(){
string s;
cin>>s,n=0,ans=0;
for(auto i:s) ++n,sum[n]=sum[n-1]+(i=='0' ? -1:1);
map<int,long long> w;
w[0]=1ll;
for(int i=1;i<=n;i++) ans=(ans+(n-i+1)*w[sum[i]])%mod,w[sum[i]]+=i+1;
cout<<ans<<'\n';
}
int T;
int main(){
ios::sync_with_stdio(0);
cin.tie(0);cout.tie(0);
cin>>T;
while(T--) solution();
return 0;
}