题解：洛谷 AT_abc441_e A B substring

【题目来源】

洛谷：AT_abc441_e [ABC441E] A > B substring - 洛谷

【题目描述】

给定一个由 A、B、C 三种字符组成的长度为 \(N\) 的字符串 \(S\)。

\(S\) 的非空连续子串共有 \(\dfrac{N(N+1)}2\) 个，请问其中包含的 A 数量多于 B 数量的子串有多少个？

请注意，即使两个子串内容相同，只要它们在 \(S\) 中取出的位置不同，就视为不同的子串。

子串 是指从 \(S\) 的开头删除 \(0\) 个或更多字符，以及从末尾删除 \(0\) 个或更多字符所得到的字符串。
例如，AB 是 ABC 的子串，但 AC 不是 ABC 的子串。

【输入】

输入通过标准输入给出，格式如下：

\(N\) \(S\)

【输出】

输出 \(S\) 的连续子串中，A 的数量多于 B 的数量的子串个数。

【输入样例】

10
ACBBCABCAB

【输出样例】

8

【核心思想】

1. 问题分析：给定长度为 \(N\) 的字符串 \(S\)（仅含 A、B、C），求所有非空连续子串中 A 数量多于 B 数量的子串个数。将 A 视为 \(+1\)、B 视为 \(-1\)、C 视为 \(0\)，则问题转化为：求满足子串和 \(> 0\) 的子串数量。这是一个前缀和 + 树状数组求逆序对问题，核心在于将子串和问题转化为前缀和的相对大小关系。

2. 算法选择：

   • 前缀和转化：设 \(pre[i]\) 为前 \(i\) 个字符的累计值（A=1, B=-1, C=0）。子串 \([l+1, r]\) 的和 = \(pre[r] - pre[l]\)。要求子串和 \(> 0\)，即 \(pre[r] > pre[l]\)（\(0 \leq l < r \leq N\)）
   • 树状数组统计：问题转化为统计满足 \(l < r\) 且 \(pre[l] < pre[r]\) 的数对个数。遍历 \(r\) 从 \(1\) 到 \(N\)，用树状数组查询已出现的前缀和中小于 \(pre[r]\) 的数量

3. 关键步骤：

   • 读入数据：\(N\) 和字符串 \(S\)
   • 字符映射：A \(\to 1\)、B \(\to -1\)、C \(\to 0\)
   • 前缀和计算：\(pre[i] = pre[i-1] + a[i]\)
   • 树状数组统计：
     • 初始化：\(add(M, 1)\)（\(pre[0] = 0\) 的情况，\(M\) 为偏移量避免负数下标）
     • 遍历 \(i\) 从 \(1\) 到 \(N\)：
       • \(ans \leftarrow ans + query(pre[i] + M - 1)\)（查询已出现的前缀和中 \(< pre[i]\) 的数量）
       • \(add(pre[i] + M, 1)\)（将当前前缀和加入树状数组）
   • 输出结果：\(ans\)

4. 时间/空间复杂度：

   • 时间复杂度：\(O(N \log N)\)，树状数组每次操作 \(O(\log N)\)
   • 空间复杂度：\(O(N)\)，前缀和数组和树状数组

5. 前缀和转化与树状数组的核心思想：

   • 子串和的前缀和表示：子串 \([l+1, r]\) 的和 = \(pre[r] - pre[l]\)，要求 \(> 0\) 即 \(pre[r] > pre[l]\)。将二维区间查询转化为一维数对统计
   • 逆序对思想的正向应用：通常逆序对统计 \(pre[l] > pre[r]\)（\(l < r\)），本题要求 \(pre[l] < pre[r]\)，即"正序对"。树状数组查询前缀和小于当前值的个数，等价于统计正序对
   • 坐标偏移处理：前缀和可能为负（B 较多时），树状数组下标从 \(1\) 开始，因此加偏移量 \(M = 500005\) 将所有值映射到正数范围
   • 初始化 \(pre[0] = 0\)：空前缀的和为 \(0\)，对应子串从 \(1\) 开始的情况。代码中通过 \(add(M, 1)\) 将 \(pre[0] = 0\) 加入树状数组
   • C 字符的零贡献：C 映射为 \(0\)，不影响前缀和，自然地被排除在 A 和 B 的计数之外
   • 适用于子串和统计、前缀和转化、逆序对/正序对类问题

【解题思路】

【算法标签】

普及 #树状数组

【代码详解】

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 500005 * 2, M = 500005 + 1;
int n, a[N];
string s;
int tr[N], ans;  // tr: 树状数组

int lowbit(int x) {
    return x & -x;  // 获取最低位的1
}

// 树状数组：在位置x加c
void add(int x, int c) {
    for (int i = x; i <= n + M; i += lowbit(i))
        tr[i] += c;
}

// 树状数组：查询前缀和[1, x]
int query(int x) {
    int res = 0;
    for (int i = x; i; i -= lowbit(i))
        res += tr[i];
    return res;
}

signed main() {
    cin >> n >> s;
    s = " " + s;  // 让字符串从1开始索引

    // 将字符转换为数字
    for (int i = 1; i <= n; i++) {
        if (s[i] == 'A') a[i] = 1;      // A记为1
        else if (s[i] == 'B') a[i] = -1; // B记为-1
    }

    // 计算前缀和
    for (int i = 1; i <= n; i++)
        a[i] += a[i-1];  // a[i]表示前i个字符的和

    add(M, 1);  // 初始化：前缀和为0的情况
    for (int i = 1; i <= n; i++) {
        // 查询前缀和小于a[i]的数量
        ans += query(a[i] + M - 1);
        // 将当前前缀和加入树状数组
        add(a[i] + M, 1);
    }

    cout << ans << endl;
    return 0;
}

【运行结果】

10
ACBBCABCAB
8