和为 K 的子数组
题目要求
编写一个函数,该函数接收两个参数:一个整数数组 nums
和一个整数 k
。你需要在这个数组中找出和为 k
的子数组的数量。子数组定义为数组中的一个连续非空序列。
解题思路
解决这个问题的关键是使用累积和的概念,以及哈希表来优化查找过程。以下是解题的步骤:
-
初始化一个哈希表,用来存储各个累积和出现的次数,以及一个变量来记录累积和的总数。在哈希表中先放入一个键值对(0:1),表示累积和为 0 的情况出现了一次,这是为了处理那些从数组第一个元素开始累加就等于 k 的子数组。
-
遍历数组
nums
,计算从第一个元素到当前元素的累积和。对于每个累积和,计算当前累积和 - k
的值,查看这个值是否已经在哈希表中存在:- 如果存在,说明找到了一个子数组的累积和等于 k(因为当前累积和减去之前某个位置的累积和等于 k,那么这两个位置之间的子数组和就是 k)。
- 将这个差值对应的次数加到计数器上,因为可能存在多个不同的子数组累积和等于这个差值。
-
接下来,将当前的累积和作为键,其出现的次数作为值,存入哈希表。如果这个累积和之前已经出现过,则增加其计数。
-
完成数组的遍历后,计数器中的值就是和为 k 的子数组的总数。
这种方法的时间复杂度为 O(n),因为只需要遍历一次数组,空间复杂度为 O(n),因为最坏情况下可能需要存储所有累积和的情况。
Golang 版本
package main
import "fmt"
func subarraySum(nums []int, k int) int {
count := 0
sum := 0
sumMap := make(map[int]int)
sumMap[0] = 1 // 初始化,累积和为0的情况出现一次
for _, num := range nums {
sum += num // 累加当前元素
if prevCount, ok := sumMap[sum-k]; ok {
count += prevCount // 如果存在sum-k的累积和,累加对应的次数
}
sumMap[sum]++ // 更新累积和出现的次数
}
return count
}
func main() {
nums := []int{1, 1, 1}
k := 2
result := subarraySum(nums, k)
fmt.Println(result) // 输出应该为2
}
这段代码定义了一个subarraySum
函数,它接收一个整数数组nums
和一个整数k
作为参数,并返回和为k
的子数组的个数。在main
函数中,我们定义了一个示例数组和k
的值,调用subarraySum
函数,并打印出结果。
C++版本
#include <iostream>
#include <vector>
#include <unordered_map>
int subarraySum(std::vector<int>& nums, int k) {
int count = 0;
int sum = 0;
std::unordered_map<int, int> sumMap;
sumMap[0] = 1; // 初始化,累积和为0的情况出现一次
for (int num : nums) {
sum += num; // 累加当前元素
if (sumMap.find(sum - k) != sumMap.end()) {
count += sumMap[sum - k]; // 如果存在sum-k的累积和,累加对应的次数
}
sumMap[sum]++; // 更新累积和出现的次数
}
return count;
}
int main() {
std::vector<int> nums = {1, 1, 1};
int k = 2;
std::cout << subarraySum(nums, k) << std::endl; // 输出应该为2
return 0;
}
这段 C++代码实现了与 Golang 版本相同的功能。它定义了一个subarraySum
函数,用于计算和为k
的子数组的数量。在main
函数中,我们创建了一个示例数组和一个k
值,调用了subarraySum
函数,并打印出了结果。
Python 版本
def subarraySum(nums, k):
count = 0
current_sum = 0
sum_dict = {0: 1}
for num in nums:
current_sum += num
if current_sum - k in sum_dict:
count += sum_dict[current_sum - k]
sum_dict[current_sum] = sum_dict.get(current_sum, 0) + 1
return count
# 测试代码
nums = [1, 1, 1]
k = 2
print(subarraySum(nums, k)) # 应该输出2
Java 版本
import java.util.HashMap;
public class Solution {
public int subarraySum(int[] nums, int k) {
int count = 0;
int sum = 0;
HashMap<Integer, Integer> sumMap = new HashMap<>();
sumMap.put(0, 1); // 初始化,累积和为0的情况出现一次
for (int num : nums) {
sum += num; // 累加当前元素
if (sumMap.containsKey(sum - k)) {
count += sumMap.get(sum - k); // 如果存在sum-k的累积和,累加对应的次数
}
sumMap.put(sum, sumMap.getOrDefault(sum, 0) + 1); // 更新累积和出现的次数
}
return count;
}
public static void main(String[] args) {
Solution solution = new Solution();
int[] nums = {1, 1, 1};
int k = 2;
System.out.println(solution.subarraySum(nums, k)); // 应该输出2
}
}
这段 Java 代码定义了一个名为Solution
的类,其中包含了subarraySum
方法,用于计算和为k
的子数组的数量。在main
方法中,我们创建了一个示例数组和一个k
值,调用了subarraySum
方法,并打印出了结果。
Rust 版本
use std::collections::HashMap; fn subarray_sum(nums: Vec<i32>, k: i32) -> i32 { let mut count = 0; let mut sum = 0; let mut sum_map: HashMap<i32, i32> = HashMap::new(); sum_map.insert(0, 1); // 初始化,累积和为0的情况出现一次 for &num in &nums { sum += num; // 累加当前元素 if let Some(&value) = sum_map.get(&(sum - k)) { count += value; // 如果存在sum-k的累积和,累加对应的次数 } *sum_map.entry(sum).or_insert(0) += 1; // 更新累积和出现的次数 } count } fn main() { let nums = vec![1, 1, 1]; let k = 2; println!("{}", subarray_sum(nums, k)); // 应该输出2 }
这段 Rust 代码实现了一个subarray_sum
函数,它接收一个整数向量nums
和一个整数k
作为参数,并返回和为k
的子数组的数量。在main
函数中,我们定义了一个示例向量和k
的值,调用了subarray_sum
函数,并打印出了结果。
总结
上述解法采用了一种称为“前缀和”的技术,结合哈希表来优化查找过程。核心思想是遍历数组,计算从开始到当前元素的累积和,并在遍历的过程中使用哈希表记录下所有出现过的累积和及其出现的次数。
具体步骤如下:
- 初始化一个计数器
count
用于记录和为k
的子数组的数量,一个变量sum
用于存储从数组开始到当前元素的累积和。 - 创建一个哈希表
sumMap
,用于存储各个累积和出现的次数,初始时存入键值对(0, 1)
,表示累积和为 0 的情况出现了一次。 - 遍历数组
nums
,对于每个元素,将其加到sum
上更新当前的累积和。 - 查看
sumMap
中是否存在键为sum - k
的项,如果存在,说明从某个位置到当前位置的子数组和为k
,将该键对应的值加到count
上。 - 更新
sumMap
,将当前的累积和sum
作为键,其出现的次数作为值。如果这个累积和之前已经出现过,则增加其计数。 - 遍历完成后,
count
变量中存储的就是和为k
的子数组的数量。
这种方法的时间复杂度为 O(n),因为它只需要遍历数组一次。空间复杂度为 O(n),因为在最坏的情况下,可能需要存储数组中所有元素的累积和。这种方法比暴力解法(双重循环遍历所有子数组)要高效得多。