和为 k 的子数组个数

解题思路

本题属于前缀和 + 哈希表的经典应用。我们需要找到所有连续子数组的和为 k 的个数。以下是解题的核心思路：

1. 前缀和思路

• 定义前缀和：prefixSum[i] 表示从数组的起始位置到第 i 个元素的累加和。
• 假设某一段子数组的和为 k，即：

• 变形得到：

这意味着，只需要找到某个前缀和 prefixSum[i]，使得它满足上述等式即可。

2. 哈希表优化

• 使用哈希表 prefixSumCount 来记录每个前缀和出现的次数。
• 初始化 prefixSumCount[0] = 1，表示前缀和为 0 出现 1 次（这种初始化用于满足从第一个元素开始的情况）。
• 遍历数组时，不断更新当前的前缀和 currentSum，并在哈希表中查询 currentSum - k 出现的次数。每次找到一个这样的 prefixSum，就意味着找到了一个满足条件的子数组。

3. 算法步骤

• 初始化前缀和计数器 prefixSumCount，并设置 prefixSumCount[0] = 1。
• 初始化 currentSum = 0 和 count = 0。
• 遍历数组：
  • 累加当前元素到 currentSum。
  • 检查 prefixSumCount[currentSum - k]，若存在，则将其次数加到 count。
  • 将 currentSum 添加到 prefixSumCount，若已存在则次数加一。
• 最终输出 count。

4. 时间复杂度分析

• 时间复杂度：O(n) — 只需一次遍历。
• 空间复杂度：O(n) — 哈希表的最大空间可能达到 n。

代码实现

Python 代码

class Solution:
    def subarraySum(self, nums, k):
        # 哈希表记录前缀和及其出现次数
        prefixSumCount = {0: 1}
        currentSum = 0
        count = 0
        
        # 遍历数组
        for num in nums:
            currentSum += num
            
            # 检查是否存在满足条件的前缀和
            if currentSum - k in prefixSumCount:
                count += prefixSumCount[currentSum - k]
            
            # 更新当前的前缀和次数
            prefixSumCount[currentSum] = prefixSumCount.get(currentSum, 0) + 1
        
        return count

# 输入输出处理
if __name__ == "__main__":
    n, k = map(int, input().strip().split())
    nums = list(map(int, input().strip().split()))
    solution = Solution()
    print(solution.subarraySum(nums, k))

Java 代码

import java.util.HashMap;
import java.util.Scanner;

public class Main {
    public class Solution {
        public int subarraySum(int[] nums, int k) {
            // 哈希表记录前缀和及其出现次数
            HashMap<Integer, Integer> prefixSumCount = new HashMap<>();
            prefixSumCount.put(0, 1);

            int currentSum = 0;
            int count = 0;

            // 遍历数组
            for (int num : nums) {
                currentSum += num;

                // 检查是否存在满足条件的前缀和
                if (prefixSumCount.containsKey(currentSum - k)) {
                    count += prefixSumCount.get(currentSum - k);
                }

                // 更新当前的前缀和次数
                prefixSumCount.put(currentSum, prefixSumCount.getOrDefault(currentSum, 0) + 1);
            }

            return count;
        }
    }

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int k = sc.nextInt();
        int[] nums = new int[n];
        for (int i = 0; i < n; i++) {
            nums[i] = sc.nextInt();
        }
        Main main = new Main();
        Solution solution = main.new Solution();
        System.out.println(solution.subarraySum(nums, k));
    }
}

C++ 代码

#include <iostream>
#include <unordered_map>
#include <vector>
using namespace std;

class Solution {
public:
    int subarraySum(vector<int>& nums, int k) {
        unordered_map<int, int> prefixSumCount;
        prefixSumCount[0] = 1; // 初始化

        int currentSum = 0;
        int count = 0;

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

            // 检查是否存在满足条件的前缀和
            if (prefixSumCount.find(currentSum - k) != prefixSumCount.end()) {
                count += prefixSumCount[currentSum - k];
            }

            // 更新前缀和次数
            prefixSumCount[currentSum]++;
        }

        return count;
    }
};

int main() {
    int n, k;
    cin >> n >> k;
    vector<int> nums(n);
    for (int i = 0; i < n; i++) {
        cin >> nums[i];
    }
    
    Solution solution;
    cout << solution.subarraySum(nums, k) << endl;
    return 0;
}

题目内容

给你一个整数数组$nums$和一个整数$k$，请你统计并输出 该数组中和为 $k$ 的子数组的个数。

子数组是数组中元素的连续非空序列。

输入描述

输入共两行。

• 第一行为两个个整数,$n，k$。

• 第二行为$n$个整数$nums_0,nums_1,...,nums_{n-1}$，数字之间以空格分隔。

输出描述

一个整数，表示答案。

样例1

输入

3 2
1 1 1

输出

2

样例2

输入

3 3
1 2 3

输出

2

提示

• $1 <= nums.length <= 2 * 10^4$
• $-1000 <= nums[i] <= 1000$
• $-10^7 <= k <= 10^7$