题目：ABC441 E：A > B substring
关键词：前缀和 / 树状数组 / 统计小于当前数的前缀
核心：把 A/B/C 转成 1/-1/0，问题变成统计子区间和大于 0 的个数。

题意

在一个只有 A, B, C 的长度为 N 的字符串中，求 A 出现次数严格大于 B 出现次数的子区间个数。

想法

注意 C 除了占用位置完全没用，我们想迅速判断一个区间合不合法应该怎么办呢？

最简单的办法：分别看 A 的数量和 B 的数量。

但是这样需要看两个指标，要分别求两个前缀和，好像也看不出来什么。

我们可以把字母转换成分数：

  • A $\rightarrow$ 1
  • B $\rightarrow$ -1
  • C $\rightarrow$ 0

这样我们就只需要找子区间和大于 0 的区间个数就好了。

利用前缀和，对于区间 $[L, R]$，合法条件是：

$$
pre[R] - pre[L - 1] > 0 \rightarrow pre[R] > pre[L - 1]
$$

也就是要统计每个数左边有几个比它小的前缀和。这是一个经典问题，可以用树状数组，也可以用归并排序。

这里用树状数组。注意前缀和可能出现负数，所以整体加一个偏移量。

实现

枚举每个前缀和 pre[i]：

  1. 查询已经出现过的、小于 pre[i] 的数量
  2. 加到答案里
  3. 把当前 pre[i] 加入树状数组

代码

代码 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
#include <bits/stdc++.h>
#define endl '\n'
#define int long long
#define lowbit(x) x & -x
const int mod = 998244353;
const int maxn = 1e6 + 5;
int n; int d[maxn];
int tr[maxn];
int pre[maxn];

void add(int pos, int x)
{
    while(pos <= maxn)
    {
        tr[pos] += x;
        pos += lowbit(pos);
    }
}

int sum(int pos)
{
    int res = 0;
    while(pos)
    {
        res += tr[pos];
        pos -= lowbit(pos);
    }
    return res;
}

void solve()
{
    std::cin >> n; char c;
    for (int i = 1; i <= n; i++) 
    {
        std::cin >> c;  
        pre[i] = pre[i - 1];
        if (c == 'A') d[i] = 1;
        if (c == 'B') d[i] = -1;
        if (c == 'C') d[i] = 0;
        pre[i] += d[i];   
    }
    for (int i = 0; i <= n; i++) pre[i] += 500001;
    int ans = 0;
    for (int i = 0; i <= n; i++)
    {
        int cur = pre[i];
        ans += sum(cur - 1);
        add(pre[i], 1);
    }
    std::cout << ans << endl;
}

signed main()
{
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int _t = 1; // std::cin >> _t;
    while (_t--) solve();
    return 0;
}