動機

我忘了prefix sum!! 同時，這題不能用sliding window

Problem

Given an array of integers nums and an integer k, return the total number of continuous subarrays whose sum equals to k.

 

Example 1:

Input: nums = [1,1,1], k = 2Output: 2

Example 2:

Input: nums = [1,2,3], k = 3Output: 2

 

Constraints:

  • 1 <= nums.length <= 2 * 104
  • -1000 <= nums[i] <= 1000
  • -107 <= k <= 107

Sol

這題不能用sliding window，因為array中有負數，不會越加越多

這樣不能調左邊的index讓window的總和縮小

所以可以用prefix sum來算區間和，之後再找對應的和

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 忘了可以用presum
        
        sums, cnt, ret = [{}, 0, 0]
        for i in range(len(nums)):
            cnt += nums[i]
            if cnt not in sums:
                sums[cnt] = [i]
            else:
                sums[cnt].append(i)
        #print(sums)
        for left in sums.keys():
            # k ==  sums[j] - sums[i]
            right = k + left
            if left == k: # 可以用在sums放來省略這一段
                ret += len(sums[left])
            if right in sums:
                for i in sums[left]:
                    tmp = sums[right]
                    ret += len(tmp)-bisect_left(tmp,i+1)
        return ret

原本是

a[a]+...+a[b] = k

變成

sum[b]-sum[a-1] = k
=>
sum[b]-k = sum[a-1]

是不是只要是連續區間和就可以用這種手段

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        ret = 0
        tbl = {0:1} # sum[i] is equal to k!!
        cnt = 0
        # sum(nums[j:i]) = sums[:i] - sum[:j] - 
        for n in nums:
            cnt += n
            if cnt-k in tbl:
                ret += tbl[cnt-k]
            tbl[cnt] = tbl.get(cnt,0)+1
        return ret

sum[i:j] = sum[:j]-sum[:i] 上面會用bisect_left，是因為要找比當前位置大的index，也就是找j，但其實可以改找i

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        ret, cnt, tbl = [0, 0 , {0:1}]
        # sum(nums[j:i]) = sums[:i] - sum[:j]
        for n in nums:
            cnt += n
            ret += tbl.get(cnt-k,0)
            tbl[cnt] = tbl.get(cnt,0)+1
        return ret