https://codeforces.com/contest/1996/problem/E
题目大意
给出 $01$ 串 $S$,长度 $\le 2 \cdot 10^5$。求其所有子串的 $01$ 个数相同的子串的个数。
简要题解
题目求 $\sum_l \sum_r \sum_x \sum_y [cnt0(x, y) == cnt1(x, y)] ,\ (1 \le l \le x \le y \le r \le n)$。
如果只有所有子串 $S(x, y)$ 的个数,那将非常好求。遇到 $0$ 就 $-1$,遇到 $1$ 就 $+1$,然后每个位置作为右端点,看前缀和里有多少个相同的值,就是有多少个合法的左端点。
子串的子串只需要进一步扩展这个想法,如果我们固定一个组合法的 $(x, y)$,显然任意 $1 \le l \le x$ 和任意 $y \le r \le n$ 都是合法的。于是我们依旧枚举 $y$,我们把所有 $+1-1$ 前缀和相同的 $x$ 对应 $l$ 的方案数记录起来即可。
复杂度
$T$:$O(n \log n)$
$S$:$O(n)$
代码实现
#include <bits/stdc++.h>
using namespace std;
int io_=[](){ ios::sync_with_stdio(false); cin.tie(nullptr); return 0; }();
using LL = long long;
using ULL = unsigned long long;
using LD = long double;
using PII = pair<int, int>;
using VI = vector<int>;
using MII = map<int, int>;
template<typename T> void cmin(T &x,const T &y) { if(y<x) x=y; }
template<typename T> void cmax(T &x,const T &y) { if(x<y) x=y; }
template<typename T> bool ckmin(T &x,const T &y) {
return y<x ? (x=y, true) : false; }
template<typename T> bool ckmax(T &x,const T &y) {
return x<y ? (x=y, true) : false; }
template<typename T> void cmin(T &x,T &y,const T &z) {// x<=y<=z
if(z<x) { y=x; x=z; } else if(z<y) y=z; }
template<typename T> void cmax(T &x,T &y,const T &z) {// x>=y>=z
if(x<z) { y=x; x=z; } else if(y<z) y=z; }
// mt19937 rnd(chrono::system_clock::now().time_since_epoch().count());
// mt19937_64 rnd_64(chrono::system_clock::now().time_since_epoch().count());
template<long long Mo=998244353> struct ModInt {
static long long MO;
static void setMo(long long mo) { MO = mo; }
long long x;
ModInt(long long x=0) : x(x){ norm(); }
friend istream &operator>>(istream& in, ModInt &B) { in>>B.x; return in; }
friend ostream &operator<<(ostream& out, const ModInt &B) {
out<<B.x; return out; }
// ModInt operator=(int x_) { x=x_; norm(); return *this; }
void norm() { x = (x%MO + MO) % MO; }
long long get() { return x; }
ModInt operator-() const { return ModInt(MO - x); }
ModInt operator+=(const ModInt &B) { x+=B.x; if(x>=MO) x-=MO; return *this; }
ModInt operator-=(const ModInt &B) { x-=B.x; if(x<0) x+=MO; return *this; }
ModInt operator*=(const ModInt &B) { x=x*B.x%MO; return *this; }
ModInt operator+(const ModInt &B) const { ModInt ans=*this; return ans+=B; }
ModInt operator-(const ModInt &B) const { ModInt ans=*this; return ans-=B; }
ModInt operator*(const ModInt &B) const { ModInt ans=*this; return ans*=B; }
ModInt operator^(long long n) const {
ModInt a=*this; ModInt ans(1);
while(n) { if(n&1) ans*=a; a*=a; n>>=1; }
return ans;
}
ModInt inv() const { return (*this)^(MO-2); } // if MO is prime
ModInt operator/=(const ModInt &B) { (*this)*=B.inv(); return *this; }
ModInt operator/(const ModInt &B) const { ModInt ans=*this; return ans/=B; }
bool operator<(const ModInt &B) const { return x<B.x; }
bool operator==(const ModInt &B) const { return x==B.x; }
bool operator!=(const ModInt &B) const { return x!=B.x; }
};
template<long long Mo> long long ModInt<Mo>::MO = Mo;
// typedef ModInt<998244353> Mint;
typedef ModInt<1'000'000'007> Mint;
/*
---------1---------2---------3---------4---------5---------6---------7---------
1234567890123456789012345678901234567890123456789012345678901234567890123456789
*/
void solve() {
string s; cin >> s;
int n = s.length();
map<int, Mint> cnt;
cnt[0] = 1;
int cur = 0;
Mint ans = 0;
for (int i = 0; i < n; i++) {
if (s[i] == '0') cur--;
else cur++;
Mint nans = cnt[cur] * Mint(n - i);
// cerr << i << ' ' << nans << endl;
ans += nans;
cnt[cur] += i + 2;
}
cout << ans << '\n';
}
int main() {
int t = 1;
cin >> t;
while (t--) {
solve();
}
return 0;
}
Next: F. Bomb - Codeforces Round 962 (Div. 3)