Leetcode/Lintcode

Leetcode 560. Subarray Sum Equals K

Description

Origin Problem

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals k.

Example 1:

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

Solution

We can use the prefix sum technique to solve the problem. Instead of calculating the sum of all possible subarrays(O(n) each), we can use the prefix sums to get the sum of subarray (i, j) in O(n). Now we just have to iterate through all (i, j) pairs and count how many sums are equal to k. The total time complexity is O(n^2).

Optimization

We can do better. Iterating through all (i, j) pairs is expensive. We can instead use a dict to store the number of each specific prefix sum. Since we want prefix[i]-prefix[j] = k, we can easily know how many js are valid, by finding the numbers of prefix sums that are k-prefix[j] in dict. To avoid counting duplicates, we calculate the prefix sums along with the counting process. The time complexity is now O(n).

Here is my code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from collections import defaultdict
class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        prefix_sum = 0
        sum_lookup = defaultdict(int) # values of dict would be initialized to 0
        sum_lookup[0] = 1
        count = 0
        
        # find all pairs of (i, j) such that (prefix_sum at i)-(prefix_sum at j) == k
        for i in range(len(nums)):
            prefix_sum += nums[i] # calculate the current prefix sum
            
            if (prefix_sum - k in sum_lookup):
                # add the numbers of 'j's which (prefix_sum at j) == (prefix_sum at i)-k
                count += sum_lookup[prefix_sum - k] 
            
            sum_lookup[prefix_sum] += 1
        return count