Leetcode 560. Subarary Sum Equals K

Problem

Leetcode 560 - Subarary Sum Equals K

Example:

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

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

Approach

Since $\sum_{x=i}^{j} a_x = \sum_{a=0}^{j} a_a - \sum_{b=0}^{i} a_b$, we can use this information to store the count of the prefix ($\sum_{b=0}^{i} a_a$).

Since $k = \sum_{x=i}^{j} a_x$, and sum of upto j is = $\sum_{a=0}^{j} a_a$, the amount of subarray that result in $k$ is equal to count of prefix ($\sum_{b=0}^{i} a_a$) upto $j$.

Thus, we can store the frequency of the prefix in the HashTable, then we can simply add the frequency of the difference between prefix and k for the count of subArray.

Visualization of the Approach:

num = [1, 1, -1, 1, 2, 5], k = 7
count = { 0 : 1 }

prefix = 0

prefix = 1
diff = prefix - k = -6
Since we don't have -6 in count, we don't increment result
count = { 0 : 1, 1: 1 }

prefix = 2
diff = prefix - k = -5
Since we don't have -5 in count, we don't increment result
count = { 0 : 1, 1 : 1, 2 : 1 }

prefix = 1
diff = prefix - k = -6
Since we don't have -6 in count, we don't increment result
count = { 0 : 1, 1 : 2, 2 : 1 }

prefix = 2
diff = prefix - k = -5
Since we don't have -5 in count, we don't increment result
count = { 0 : 1, 1 : 2, 2 : 2 }

prefix = 4
diff = prefix - k = -3
Since we don't have -3 in count, we don't increment result
count = { 0 : 1, 1: 2, 2 : 2, 4 : 1 }

prefix = 9
diff = prefix - k = 2
Since there is 2 in count, we increment result by 2: result = 2
count = { 0 : 1, 1 : 2, 2 : 2, 4 : 1, 9 : 1 }

Thus, result is 2

Here is the Python code for the solution:

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        res = 0
        prefix = 0
        # initialize the count of 0 to 1 since even before we start, the summation of prefix is 0
        count = { 0 : 1 }

        for num in nums:
            prefix += num
            diff = prefix - k

            res += count.get(diff, 0)
            count[prefix] = count.get(prefix, 0) + 1 
        
        return res        

Time Complexity and Space Complexity

Time Complexity: $O(n)$

Space Complexity: $O(n)$