Codeforces Round 962 (Div. 3) E(思维 + 计数)

题面

给定一个长度为 $n(1\le n\le 2\times {10}^5)$ 的 $01$串 $S$，问：对于每一对 $l,r(1\le l\le r\le n)$，有多少对 $(x,y)$ 满足：$l\le x\le y\le r$ 且 $S_{x\sim y}$ 中字符 $1$ 的个数和字符 $0$ 的个数相同。

题解

实际上，我们只需要找出每一对 $(x,y)$，满足 $S_{x\sim y}$ 中字符 $0$ 和 $1$ 数量相同，看它对应的数对 $(l,r)$ 的数量是多少，将所有的数对 $(l,r)$ 的数量求和即可。

如果 $(x,y)$ 确定下来了，那么 $l$ 的可行范围在 $[1,x]$，$r$ 的可行范围在 $[y,n]$，即：数对 $(l,r)$ 的数量为 $x\times (n-y+1)$。

我们可以把字符 $0$ 看作数字 $-1$，把字符 $1$ 看作数字 $1$，然后做前缀和，将前 $i$ 项的和记作 $pr{e}_i$。如果 $pr{e}_i=pr{e}_j$，那么 $S_{i\sim j}$ 中的字符 $0$ 和字符 $1$ 数量相同。

可以考虑枚举每个位置作为 $r$，然后将其贡献累加起来，对每一个 $pr{e}_i$，记录其对应的 $i$ 的和，这样，位置 $j$ 作为 $r$，对答案的贡献就是 $su{m}_{pr{e}_j}\times (n-j+1)$。

时间复杂度为 $\mathcal{O}(nlogn)$

点击查看代码

#include <bits/stdc++.h>

using i64 = int64_t;

void solve() {
  std::string s;
  std::cin >> s;

  const int N = s.size();
  std::vector<int> pref(N + 1);
  for (int i = 0; i < N; i++) {
    pref[i + 1] = pref[i] + (s[i] == '0' ? -1 : 1);
  }

  std::map<int, i64> f;
  f[0] = 1;

  Z ans = 0;
  for (int i = 1; i <= N; i++) {
    // Z 是取模类的名字
    ans += Z(N - i + 1) * f[pref[i]];
    f[pref[i]] += i + 1;
  }
  std::cout << ans << "\n";
}

int main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr);

  int T;
  std::cin >> T;

  while (T--) {
    solve();
  }

  return 0;
}