和为 K 的子数组

题意

给定整数数组 nums 和整数 k，统计和为 k 的连续子数组的个数。

思路

暴力 O(n²) 枚举所有子数组。优化：前缀和 s[i] 表示 [0..i) 的和，子数组 [l..r] 的和为 s[r+1] - s[l]。遍历时用哈希表记录每个前缀和出现的次数，对当前位置 sum，查 sum - k 的出现次数即为以当前位置结尾的合法子数组个数。

注意：先查询后插入，且初始插入 {0: 1} 表示空前缀。

代码

cpp

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/**
 * Author by Rainboy
 */
// main.cpp：前缀和 + 哈希表统计前缀出现次数，O(n)。
#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    int subarraySum(vector<int> &nums, int k) {
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        int sum = 0, ans = 0;
        for (int x : nums) {
            sum += x;
            auto it = cnt.find(sum - k);
            if (it != cnt.end())
                ans += it->second;
            cnt[sum]++;
        }
        return ans;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, k;
    cin >> n >> k;
    vector<int> a(n);
    for (int &x : a)
        cin >> x;
    cout << Solution().subarraySum(a, k) << '\n';
    return 0;
}

python

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#!/usr/bin/env python3
from typing import List
from collections import defaultdict


class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        cnt = defaultdict(int)
        cnt[0] = 1
        s = ans = 0
        for x in nums:
            s += x
            ans += cnt[s - k]
            cnt[s] += 1
        return ans


def main() -> None:
    n, k = map(int, input().split())
    nums = list(map(int, input().split()))
    print(Solution().subarraySum(nums, k))


if __name__ == "__main__":
    main()

复杂度

时间复杂度：O(n)，每个元素处理一次。
空间复杂度：O(n)，哈希表最多存 n 个前缀和。

总结

前缀和配合哈希表是子数组统计问题的标准模型。与两数之和的配对计数本质相同：固定右端点，查历史信息的数量。有负数时双指针失效，但前缀和哈希表仍然正确。

目录

题意

思路

代码

复杂度

总结