#P4009. 和为K的子数组
-
ID: 2216
Tried: 173
Accepted: 77
Difficulty: 4
和为K的子数组
题目内容
给你一个整数数组nums和一个整数k，请你统计并输出 该数组中和为 k 的子数组的个数。
子数组是数组中元素的连续非空序列。
输入描述
输入共两行。
-
第一行为两个个整数,n，k。
-
第二行为n个整数nums0,nums1,...,numsn−1，数字之间以空格分隔。
输出描述
一个整数，表示答案。
样例1
输入
3 2
1 1 1
输出
2
样例2
输入
3 3
1 2 3
输出
2
提示
- 1<=nums.length<=2∗104
- −1000<=nums[i]<=1000
- −107<=k<=107
和为 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;
}