Subarray Sum Equals k

ArrayPrefix SumHash Table

# Solution

The brute force solution would:

  • enumerate subarrays in a nested for loop (O(n2)O(n^2))
  • calculate sum in a for loop (O(n)O(n))
  • update if current subarray has largest sum (O(1)O(1))

So in total takes O(n3)O(n^3) time.

# Squared Time (TLE)

The following two O(n2)O(n^2) time solutions would TLE in Python but might pass in other languages, such as C++ and Java.


# Prefix Sum

Calculate prefix sum in O(n)O(n) time beforehand, so that csum can be obtained in O(1)O(1) time.


time: O(n2)O(n^2)
space: O(n)O(n)

def subarraySum(self, nums: List[int], k: int) -> int:
    n = len(nums)
    cnt = 0

    prefix_sum = [0 for _ in range(n)]
    prefix_sum[0] = nums[0]
    for i in range(1,n):
        prefix_sum[i] = prefix_sum[i-1] + nums[i]

    for i in range(n):
        for j in range(i,n):
            if i==0:
                csum = prefix_sum[j]
                csum = prefix_sum[j] - prefix_sum[i-1]
            if csum==k:
                cnt += 1
    return cnt

# Cumulative Sum

Calculate csum cumulatively.


time: O(n2)O(n^2)
space: O(1)O(1)

def subarraySum(self, nums: List[int], k: int) -> int:
    n = len(nums)
    cnt = 0
    for i in range(n):
        csum = 0
        for j in range(i,n):
            csum += nums[j]
            if csum==k:
                cnt += 1
    return cnt


# Linear Time: HashMap

If nums[i] summing to nums[j] equals k, prefix_sum[i] == prefix_sum[j] - k, then we just need to find if satisfying prefix_sum[i] has occurred. If so, increase count by occurrences of prefix_sum[i]. Use a HashMap to store mappings from prefix sum to occurrence.

The following three solutions are all from this post (opens new window). The last two using Counter are a bit of overhaul and Vanilla HashMap is good enough.


time: O(n)O(n)
space: O(n)O(n)

# Vanilla HashMap

def subarraySum(self, nums, k):
    hm, csum, res = {0: 1}, 0, 0
    for num in nums:
        csum += num
        res += hm.get(csum - k, 0)
        hm[csum] = hm.get(csum, 0) + 1
    return res
Counter & Accumulate

# Counter

Use Counter from collections as a shorthand.

def subarraySum(self, A, K):
    count = collections.Counter()
    count[0] = 1
    ans = su = 0
    for x in A:
        su += x
        ans += count[su-K]
        count[su] += 1
    return ans

# Accumulate

from collections import Counter
from itertools import accumulate

def subarraySum(nums, k):
    total = 0
    count = Counter()
    count[0] += 1
    for acc in accumulate(nums):
        total += count[acc - k]
        count[acc] += 1
    return total