python前缀_【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