【python刷题】前缀和

求连续子数组的和

def subarraySum(nums):
    preSum = [0 for _ in range(len(nums)+1)]
    for i in range(len(nums)):
        preSum[i+1] = preSum[i] + nums[i]
    return preSum

nums = [3,5,2,-2,4,1]
res = subarraySum(nums)
print(res)

接下来我们要求连续子数组的和只需要利用：preSum[j+1]-preSum[i]

leetcode 560 和为K的子数组

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        preSum = [0 for _ in range(len(nums)+1)]
        for i in range(len(nums)):
            preSum[i+1] = preSum[i] + nums[i]
        print(preSum)
        res = 0
        for i in range(1,len(nums)+1):
            for j in range(i):
                if preSum[i] - preSum[j] == k:
                    res += 1
        return res  

优化：我直接记录下有几个sum[j]和sum[i]-k相等，直接更新结果，就避免了内层的 for 循环。我们可以用哈希表，在记录前缀和的同时记录该前缀和出现的次数。

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        dic = {}
        acc, res = 0, 0
        for num in nums:
            acc += num
            if acc == k: res += 1
            if acc - k in dic: res += dic[acc - k]
            dic[acc] = dic.get(acc, 0) + 1
        return res