560. Subarray Sum Equals K

Subarray Sum Equals K - LeetCode

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 = 2 Output: 2 Example 2: Input: nums = [1,2,3], k = 3 Output: 2

Constraints:

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

  • code TLE $$O(n^2)$$
class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        c = 0
        memo = {}
        
        memo[-1] = 0
        for i in range(len(nums)):
            memo[i] = memo[i-1] + nums[i]
        for i in range(-1, len(nums)):
            for j in range(i+1, len(nums)):
                if (memo[j] - memo[i]) == k: c += 1
        return c
  • code O(n)
class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        res = 0
        memo = defaultdict(int)
        memo[0] = 1 # sum 0. e.g. k = 7, num = 7
        
        csum = 0
        for v in nums:
            csum += v
            if csum - k in memo:
                res += memo[csum - k]
            memo[csum] += 1 # can't before if. e.g. [1], 0
        return res