5

560. 和为 k 的子数组

 2 years ago
source link: https://sexywp.com/560-subarray-sum-equals-k.htm
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

560. 和为 k 的子数组

这个题目意思比较简单，就是求一个数组中，有多少个和为 k 的子数组。输入数组的规模，最少一个元素，最多 20000 个元素。

例如：[1,2,3] k=3，返回 2，因为和为 3 的子数组有两个 [1,2] 和 [3]。所谓子数组，有两个特点，一是连续，另一个是有序。

解这道题目的最直观的办法是遍历。因为一个子数组，比如有一个开始的下标，一个结束的下标，且 0 <= start < n，0 <= end < n，且 start <= end，只要分别遍历就可以了，遍历了 start 和 end，还要求和，如果处理不好的话，会写出 O(n^3) 的算法，因为对 start 到 end 中间的数字求和复杂度也是 O(n) 的。其实，我们可以用一个变量记住 [i, j] 的和，[i, j + 1] 的和就是 [i, j] + nums[j+1]，这样整个算法就可以降到 O(n^2) 的复杂度。

还有一个办法，也可以避免求和的计算，就是使用前缀和。设 prefix[i] = sum(0, i)，prefix[j] = sum(0, j)，则 sum(i, j) = prefix[j] – prefix[i]，这样，只要事先准备好 prefix 数组，求 sum(i, j) 的时候，O(1) 就可以得到。

不过，看题目的规模达到 2 万，则 O(n^2) 最坏有 4 亿次计算，也一定会超时了。

怎么优化呢？我们要求的是 prefix[j] – prefix[i] = k，其中 i < j，我们不难发现，如果 prefix[i] = prefix[j] – k，则，prefix[j] – (prefix[j] – k) = k，所以，当我们在求前缀和的时候，看看已经出现的前缀和，等于 prefix[j] – k 的值有几个，就可以知道满足题目要求的子数组数量。

由此写出算法：

def subarraySum(self, nums: List[int], k: int) -> int:
counter = Counter([0])
res, prefix, n = 0, 0, len(nums)
for i in range(n):
prefix += nums[i]
res += counter[prefix - k]
counter[prefix] += 1
return res
def subarraySum(self, nums: List[int], k: int) -> int:
    counter = Counter([0])
    res, prefix, n = 0, 0, len(nums)

    for i in range(n):
        prefix += nums[i]
        res += counter[prefix - k]
        counter[prefix] += 1

    return res

这个算法的时间和空间复杂度都是 O(n)。可以比较有效的求出答案。


report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK