NekoNowcoder

Appearance

数据结构 - 题解 ​

标签与难度 ​

标签: 哈希表, 前缀和, 字符串, 计数, 思维, 状态压缩 难度: 1600

题目大意喵~ ​

你好呀，未来的算法大师！这道题是这样的喵：

我们拿到一个只包含字符 '0', '1', '2' 的字符串，长度是 nn。 我们的任务是，找出这个字符串里有多少个子串，满足一个特殊的条件：在这个子串里，'0', '1', '2' 这三个字符的数量是完全相等的！

比如说，对于字符串 "01210"：

  • 子串 "012" 是一个解，因为它有1个'0', 1个'1', 1个'2'。
  • 子串 "1210" 不是解，因为它有1个'0', 2个'1', 1个'2'。

你需要计算出所有满足条件的子串的总数，然后告诉我就好啦，喵~

解题思路分析 ​

这道题如果直接去检查每一个子串，那可就太慢了呀！一个长度为 nn 的字符串，子串的数量可是有 O(n2)O(n^2) 这么多的，再对每个子串计数，时间复杂度就更高了，对于 nn 这么大的情况，肯定会超时的说。所以，我们需要一个更聪明的办法，喵！

这种关于子串计数的问题，一个非常强大的武器就是前缀和！让我来带你一步步推导吧~

1. 从条件出发 ​

我们的目标是找到一个子串 s[j...i]，使得它包含的 '0', '1', '2' 数量相等。 我们用 count0(k), count1(k), count2(k) 分别表示在字符串的前 kk 个字符（也就是 s[0...k-1]）中 '0', '1', '2' 的数量。

那么，对于子串 s[j...i]，其中各个字符的数量就可以表示为：

  • '0' 的数量: count0(i+1) - count0(j)
  • '1' 的数量: count1(i+1) - count1(j)
  • '2' 的数量: count2(i+1) - count2(j)

我们的条件是这三个数量相等，也就是： count0(i+1) - count0(j) == count1(i+1) - count1(j)count1(i+1) - count1(j) == count2(i+1) - count2(j)

2. 神奇的变形 ​

这两个等式看起来还是有点复杂，我们来施展一下数学魔法，把它们变个形！我们把和下标 i 相关的项移到一边，和下标 j 相关的项移到另一边：

第一个等式变形后得到： count0(i+1) - count1(i+1) == count0(j) - count1(j)

第二个等式变形后得到： count1(i+1) - count2(i+1) == count1(j) - count2(j)

哇！发现了吗？这个变形揭示了一个惊人的秘密！

一个子串 s[j...i] 满足条件，当且仅当 "前缀 i 的状态" 和 "前缀 j-1 的状态" 是相同的！

这里的 "状态" 是什么呢？就是由两个差值构成的一个组合： 状态 (k) = (count0(k) - count1(k), count1(k) - count2(k))

3. 转化为计数问题 ​

现在，问题就转化成了一个我们更熟悉的样子： 我们从头到尾遍历字符串，计算到每个位置 i 为止的前缀状态 (count0(i+1) - count1(i+1), count1(i+1) - count2(i+1))。 对于每个位置 i，我们想知道，在它之前有多少个位置 j-1，其前缀状态和当前位置 i 的前缀状态是完全一样的。

这不就是一个经典的计数问题嘛！我们可以用一个哈希表（在C++里是 std::map 或者 std::unordered_map）来解决。

4. 算法流程 ​

  1. 创建一个哈希表 state_counts，用来存储每个状态 (差值1, 差值2) 已经出现过的次数。
  2. 在开始之前，一个空的前缀（可以看作是第-1个位置）是存在的，它的 '0', '1', '2' 数量都是0，所以差值是 (0, 0)。我们初始化 state_counts[{0, 0}] = 1。这代表了我们有一个“起点”。
  3. 初始化总答案 ans = 0，以及三个字符的计数器 c0 = 0, c1 = 0, c2 = 0。
  4. 遍历字符串的每个字符 s[i]： a. 根据 s[i] 的值，更新 c0, c1, c2。 b. 计算当前的前缀状态 currentState = (c0 - c1, c1 - c2)。 c. 在哈希表中查找 currentState 已经出现过多少次了。假设是 k 次。这 k 次就对应了 k 个合法的子串起始点 j（s[j...i]）。所以，我们将答案 ans 加上 k。 d. 然后，将当前这个状态 currentState 的出现次数在哈希表中加一，因为它也成为了一个新的历史状态，可以作为未来子串的起点了。
  5. 遍历结束后，ans 就是我们要求的总数啦！

举个栗子，s = "012"：

  • 开始: ans=0, c0=0, c1=0, c2=0。map里有 {(0,0): 1}。
  • i=0, s[0]='0':
    • c0=1, c1=0, c2=0。
    • 状态是 (1-0, 0-0) = (1,0)。
    • map里 (1,0) 出现0次，ans += 0。
    • map 更新为 {(0,0):1, (1,0):1}。
  • i=1, s[1]='1':
    • c0=1, c1=1, c2=0。
    • 状态是 (1-1, 1-0) = (0,1)。
    • map里 (0,1) 出现0次，ans += 0。
    • map 更新为 {(0,0):1, (1,0):1, (0,1):1}。
  • i=2, s[2]='2':
    • c0=1, c1=1, c2=1。
    • 状态是 (1-1, 1-1) = (0,0)。
    • map里 (0,0) 出现1次（就是最开始的空前缀），ans += 1。
    • map 更新为 {(0,0):2, (1,0):1, (0,1):1}。

最后 ans 是 1，正确！这个方法是不是又快又优雅呢，喵~

代码实现 ​

这是我根据上面的思路，精心为你准备的一份代码哦！注释很详细，希望能帮到你，喵~

cpp
#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <utility> // for std::pair

// 使用一个函数来处理单次测试，这样main函数会很整洁呐
void solve() {
    int n;
    std::cin >> n;
    std::string s;
    std::cin >> s;

    // 使用 long long 来存储答案，因为结果可能会很大哦！
    long long ans = 0;

    // c0, c1, c2 分别记录从字符串开头到当前位置'0', '1', '2'的数量
    int c0 = 0, c1 = 0, c2 = 0;

    // map<pair<int, int>, int> 用来存储状态出现的次数
    // Key: pair<int, int> 表示状态 (count0 - count1, count1 - count2)
    // Value: int 表示这个状态已经出现过的次数
    std::map<std::pair<int, int>, int> state_counts;

    // 初始化：空前缀的状态是 (0, 0)，它出现了一次
    // 这是所有子串的“共同起点”
    state_counts[{0, 0}] = 1;

    // 遍历整个字符串
    for (char ch : s) {
        // 1. 更新当前前缀的字符计数
        if (ch == '0') {
            c0++;
        } else if (ch == '1') {
            c1++;
        } else { // ch == '2'
            c2++;
        }

        // 2. 计算当前前缀的状态
        std::pair<int, int> current_state = {c0 - c1, c1 - c2};

        // 3. 查找这个状态之前出现过多少次
        // 如果map中存在这个状态，就意味着我们找到了若干个合法的子串起始点
        if (state_counts.count(current_state)) {
            ans += state_counts[current_state];
        }

        // 4. 将当前状态的出现次数加一，为后面的计算做准备
        state_counts[current_state]++;
    }

    std::cout << ans << std::endl;
}

int main() {
    // 加速输入输出，让程序跑得更快一点~
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    int t;
    std::cin >> t;
    while (t--) {
        solve();
    }

    return 0;
}

复杂度分析 ​

  • 时间复杂度: O(Nlog⁡N)O(N \log N) 我们对长度为 NN 的字符串只遍历了一次。在循环的每一步中，我们都对 std::map 进行了操作（查找和插入）。std::map 的操作时间复杂度是对数级别的，即 O(log⁡K)O(\log K)，其中 KK 是 map 中不同状态的数量。在最坏的情况下，KK 可能接近 NN，所以总的时间复杂度是 O(Nlog⁡N)O(N \log N)。如果使用 std::unordered_map，平均时间复杂度可以优化到 O(N)O(N) 哟！

  • 空间复杂度: O(N)O(N)state_counts 这个哈希表最多会存储 N+1N+1 个不同的状态（包括初始的空前缀）。每个状态由两个整数和一个计数器构成。因此，占用的空间与字符串长度 NN 成正比，空间复杂度是 O(N)O(N) 的说。

知识点总结 ​

这道题是前缀和思想的一个非常漂亮的应用，喵~ 通过它我们可以学到：

  1. 前缀和/差分思想: 遇到子数组/子串相关的问题，特别是和区间内的元素总和、数量、属性有关时，优先考虑前缀和。它可以将对区间的查询，转化为对两个前缀端点的查询。
  2. 状态抽象与哈希: 将一个看似复杂的条件（三个数相等）通过数学变换，抽象成一个简单的、可哈希的“状态”（比如一个 pair）。这是解决很多计数问题的关键一步。
  3. 哈希表 (std::map): 当你需要快速地记录和查询某个“状态”或“元素”的出现次数时，哈希表是你的不二之选。它把查找操作从线性时间优化到了对数或常数时间。

希望这篇题解能让你有所收获，继续加油哦，你超棒的！喵~ >w<