前缀和

十分邪恶的区间子区间问题,还是先理解题意。

给定一个 $\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;
}
最后修改:2024 年 09 月 29 日
如果觉得我的文章对你有用,请随意赞赏