LeetCode 560 和为 K 的子数组

85 阅读3分钟

leetcode.cn/problems/su…

image.png

解法一:前缀和

func subarraySum(nums []int, k int) int {
        res := 0
        // 构造前缀和数组
        preSumList := make([]int, len(nums)+1) // 表示前n个数的和,多一个元素表示0个数的和
        preSumList[0] = 0
        for i, v := range nums{
            preSumList[i+1] = preSumList[i] + v
        }
        // 固定左边界,枚举右边界,暴力穷举每次窗口内的子数组和
        // 窗口的范围为[left, right],左闭右闭
    	for left := 1; left <= len(nums); left++{
            for right := left; right<=len(nums);right++{ // 右边界每次从left开始,是考虑只有一个值的子数组即nums[left]
                // [left, right]之间的子数组和等于[0, right]的子数组和减去 [0, left-1]的子数组和
                if preSumList[right] - preSumList[left-1] == k{ 
                    res++
                    // 这里不能break,由于数组存在负数,因此找到一个子数组和符合答案后,还应该继续遍历,因为可能后半截的和为0,同个左边界但存在多个答案
                }
            }
        }
    	return res
    }

解法二:前缀和+哈希表

前缀和特性 image.png 题目要求找到所有子数组的和等于 k,即满足: preSum[r+1] − preSum[l] = k

推导可得,当计算到第 i 个前缀和 preSum[i] 时,若存在某个 l 使得 preSum[l] = preSum[i] - k,则说明从 l 到 i-1 的子数组(即 nums[l..i-1])的和为 k

func subarraySum(nums []int, k int) int {
    n := len(nums)
    // preSum[i] = nums[0] + nums[1] + ... + nums[i-1](即前 i个元素的和)
	preSum := make([]int, n+1)
    // 长度+1是为了表示前0个数的和,显然前0个数的和为0
	preSum[0] = 0
	// 前缀和到该前缀和出现次数的映射,方便快速查找所需的前缀和个数
	count := make(map[int]int)
	count[0] = 1
	// 记录和为 k 的子数组个数
	res := 0

	// 计算 nums 的前缀和
	for i := 1; i <= n; i++ {
		preSum[i] = preSum[i-1] + nums[i-1]
		need := preSum[i] - k // 需要查找的前缀和值,若存在 preSum[l] = need,则 sum(l, i-1) = k
		if val, ok := count[need]; ok {
            // 说明存在子数组[l, ... , i-1]的和为 k,累加其出现次数
			res += val
		}
		// 将当前的前缀和存入哈希表
		count[preSum[i]]++
	}
	return res
}

这一思路将时间复杂度从暴力解法的时间复杂度 O(n^2) 优化到 O(n),是利用空间换时间的经典思想。

优化空间复杂度

实际上,不需要关心具体是哪2个下标之间的子数组和,只关心符合答案的子数组个数,因此不需要整个前缀和数组,只需要记录某个前缀和的出现次数,从而优化时空复杂度

    func subarraySum(nums []int, k int) int {
        res := 0
        memo := make(map[int]int) // {子数组前缀和: 出现次数}
        memo[0] = 1 // 前缀和等于0必有一种情况,就是0个数
        preSum := 0
        for i := 0; i<len(nums); i++{
            preSum += nums[i] // 每个下标i代表前i+1个元素的前缀和
            // 根据当前“前缀和”,在 map 中寻找「与之相减 == k」的历史前缀和
            // 存在说明这个之前出现的前缀和,满足「当前前缀和 - 该前缀和 == k」,差分出一个子数组,它出现的次数累加给res
            if _, ok:=memo[preSum-k]; ok{ 
                res += memo[preSum-k]
            }
            memo[preSum]++
        }
    	return res
    }