“计数”问题

MrHe··3 min read

别看“计数”这两个字简单，以为是小学数学题。在算法世界里，它可是内功心法的试金石。从一个菜鸟到一个高手，往往就体现在对同一个计数问题，能给出几种不同层次的解法。

今天，咱们就拿一个经典的计数问题开刀，层层深入，感受一下思维提升的快感。

问题：和为 K 的子数组 ( LeetCode 560 )

给你一个整数数组 nums 和一个整数 k，请你统计并返回该数组中和为 k 的 连续子数组 的个数。

举个例子： nums = [1, 1, 1], k = 2 输出：2 解释：[1, 1] 和 [1, 1] 这两个子数组的和都是 2。

再来一个： nums = [1, 2, 3], k = 3 输出：2 解释：[1, 2] 和 [3] 这两个子数组的和都是 3。

解法一：暴力枚举，最朴素的浪漫

拿到任何题，先别想什么骚操作、神优化。脑子里第一个蹦出来的，最直观、最暴力的方法是啥？就是它了！先把它实现出来，保个底。

对于这道题，啥叫子数组？就是从 i 位置开始，到 j 位置结束（i <= j）的一段连续元素。那我们把所有的子数组都找出来，算一下它们的和，看看是不是等于 k，不就完事了？

思路：

  1. 写两层循环。第一层 i 从 0 到 n-1，枚举子数组的起点。

  2. 第二层 j 从 i 到 n-1，枚举子数组的终点。

  3. 这样，[i, j] 就构成了一个子数组。

  4. 再来一层循环，从 i 到 j 求和，判断和是否等于 k。

等等，第三层循环是不是有点蠢？我完全可以在 j 移动的时候，动态计算 [i, j] 的和。

改进后的暴力思路：

  1. 第一层循环 i，枚举起点。

  2. 第二层循环 j，从 i 开始，枚举终点。在循环 j 的同时，维护一个 sum 变量，表示 nums[i...j] 的和。

  3. 每次 j 往后走一步，sum 就加上 nums[j]。然后判断 sum 是否等于 k，如果是，计数器加一。

这思路清晰明了，直接上代码。

// 解法一：暴力枚举
public int subarraySum_1(int[] nums, int k) {
    int count = 0;
    // 枚举起点
    for (int start = 0; start < nums.length; start++) {
        int sum = 0;
        // 从起点开始，向后枚举终点
        for (int end = start; end < nums.length; end++) {
            // 计算从 start 到 end 的和
            sum += nums[end];
            // 如果和等于 k，计数器加一
            if (sum == k) {
                count++;
            }
        }
    }
    return count;
}

复杂度分析：

  • 时间复杂度：两层 for 循环，妥妥的 O(N²)。

  • 空间复杂度：O(1)，没用啥额外空间。

面试官看到这个解法，会点点头，心想：“嗯，这小伙子基础还行，题意理解没问题。但如果就这点东西，那今天就到这了。” 所以，咱们得继续往下走。

解法二：前缀和，预处理的威力

O(N²) 的瓶颈在哪？在于我们对于每一个起点 start，都重新向后累加求和。这里面有大量的重复计算。比如，我们算了 [0, 2] 的和，接着又算 [0, 3] 的和，其实后者就是在前者的基础上加上 nums[3]，我们的暴力解确实是这么做的。

但问题是，当我们计算 [1, 3] 的和时，又得从 nums[1] 开始重新加一遍。有没有办法，让我能在 O(1) 的时间内，瞬间知道任意一个子数组 [i, j] 的和呢？

当然有！这就是 前缀和 数组。

preSum[i] 表示 nums[0...i-1] 的和。我们先花 O(N) 的时间把这个数组算出来。 那么，nums[i...j] 的和是多少呢？ 它等于 (nums[0...j] 的和) - (nums[0...i-1] 的和)，也就是 preSum[j+1] - preSum[i]。

看，通过预处理，我们把求任意子数组和这个操作，从 O(N) 变成了 O(1)。 那整体思路就变成：

  1. 创建一个前缀和数组 preSum，长度为 n+1。preSum[i] 存储 nums[0...i-1] 的和。

  2. 两层循环枚举起点 i 和终点 j。

  3. 通过 preSum[j+1] - preSum[i] 直接计算出 nums[i...j] 的和，判断是否等于 k。

// 解法二：前缀和
public int subarraySum_2(int[] nums, int k) {
    int n = nums.length;
    // 构造前缀和数组
    // preSum[i] 表示 nums[0...i-1] 的和
    int[] preSum = new int[n + 1];
    preSum[0] = 0;
    for (int i = 0; i < n; i++) {
        preSum[i + 1] = preSum[i] + nums[i];
    }

    int count = 0;
    // 再次枚举起点和终点
    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            // sum of nums[i...j] is preSum[j+1] - preSum[i]
            if (preSum[j + 1] - preSum[i] == k) {
                count++;
            }
        }
    }
    return count;
}

复杂度分析：

  • 时间复杂度：构造前缀和数组是 O(N)，后面的两层循环还是 O(N²)。所以总的还是 O(N²)。

  • 空间复杂度：O(N)，因为我们创建了一个前缀和数组。

面试官看到这，会觉得：“哦？还知道前缀和，不错。虽然时间复杂度没变，但思路已经上了一个台阶，知道用空间换取局部操作的时间。但……还能不能再快点？”

解法三：前缀和 + 哈希表，终极优化

O(N²) 还是太慢，我们的目标是星辰大海——O(N)！

既然用了前缀和，我们的目标就变成了找有多少个 (i, j) 对，满足 preSum[j+1] - preSum[i] == k。

我们把这个公式变个形：preSum[i] = preSum[j+1] - k。

这个公式是啥意思？ 当我们在 j 位置，算出了 preSum[j+1] 之后，我们不再是去回头找从 0 到 j 所有的 i 来进行配对。而是反过来想：要满足条件，我需要一个什么样的 preSum[i]？我需要一个值为 preSum[j+1] - k 的 preSum[i]。

那么问题就转化为：当我的指针走到 j 时，在我之前（i <= j），出现过多少个前缀和，它的值恰好是 preSum[j+1] - k ？

这个问题，是不是听起来就很 “哈希表” ？

我们可以用一个 HashMap，key 存前缀和的值，value 存这个前缀和出现过的次数。

思路：

  1. 创建一个 HashMap<Integer, Integer> map，用来存 <前缀和, 出现次数>。

  2. 初始化 map.put(0, 1)。这是个关键！前缀和为 0 出现 1 次，代表一个元素都不取的时候。这可以完美处理那些从数组开头就满足条件的子数组。比如 nums = [3, 4, ...], k = 3，当走到 nums[0] 时，当前前缀和是 3，我们需要找 3 - 3 = 0 的前缀和，map 里正好有，于是 count 加 1，找到了子数组 [3]。

  3. 遍历 nums 数组，维护一个当前的前缀和 currentSum。

  4. 在每个位置 i，currentSum 就是 nums[0...i] 的和。

  5. 我们想找的，是 currentSum - k 这个值在 map 里出现过多少次。把它累加到结果 count 中。

  6. 别忘了，把自己当前这个 currentSum 也更新到 map 里，给后面的人用。

整个过程只需要遍历一次数组。

// 解法三：前缀和 + 哈希表
public int subarraySum_3(int[] nums, int k) {
    // key: 前缀和, value: 该前缀和出现的次数
    HashMap<Integer, Integer> preSumCount = new HashMap<>();
    // base case：前缀和为0的出现过1次（一个数都不选）
    preSumCount.put(0, 1);

    int count = 0;
    int currentSum = 0;
    for (int num : nums) {
        currentSum += num;

        // 我们要找的目标前缀和
        int target = currentSum - k;

        // 如果map中存在这个目标前缀和，说明从那个位置到当前位置的子数组和为 k
        if (preSumCount.containsKey(target)) {
            count += preSumCount.get(target);
        }

        // 把当前的前缀和也存入map，方便后续的计算
        preSumCount.put(currentSum, preSumCount.getOrDefault(currentSum, 0) + 1);
    }
    return count;
}

复杂度分析：

  • 时间复杂度：我们只遍历了一遍数组，HashMap 的操作近似 O(1)，所以总的是 O(N)。

  • 空间复杂度：在最坏情况下，每个前缀和都不同，HashMap 需要存储 N 个键值对，所以是 O(N)。

这才是面试官最想听到的答案。它完美地展示了你是如何一步步分析瓶颈，并利用合适的数据结构（从数组到哈希表）来优化算法的。这种“空间换时间”的思想，是算法优化的核心。

总结一下

你看，同样一个“计数”问题：

  1. 暴力解：O(N²)，是基本功，也是思考的起点。

  2. 前缀和：O(N²)，引入了预处理思想，虽然没降低最终复杂度，但优化了核心操作。

  3. 前缀和 + 哈希表：O(N)，通过巧妙的公式变形，把问题转化为查找问题，并用哈希表这个大杀器一击致命。

  • #4occ6k6h5pww4ocd6zeu6aky