[Leetcode]560. Subarray Sum Equals K

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.
Example 1:
Input:nums = [1,1,1], k = 2 Output: 2

The brute force approaching would be checking each substrings fo the input leading to O(n^2) time complexity. To improve it, we may observe that, taking arr=[1,2,3] as an example, we would exhaust 1,1+2,1+2+3,2,2+3,3, where repeated computations happen. If we store the pre-sums along our scanning, we would be able to solve it with one-pass.

Logic

As shown above, we maintain an array sum to record the sum of arr[0...j]. Our strategy is to scan the array, count the pre-sums of arr[0...j] and find the number of qualified substrings arr[i...j] that ends with arr[j].

To find the substrings arr[i...j] that are qualified, i.e., sum[i...j]==target, we need to find i that sum[0...i]==sum[0...j]-target. To achieve this in O(1), we use one hashmap dic to record the frequencies of each pre-sum.

Code

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        dic = collections.defaultdict(int)
        dic[0] = 1
        total = ret = 0
        for e in nums:
            total += e
            ret += dic[total - k]
            dic[total] += 1
        return ret