560. Subarray Sum Equals K

Problem Statement

You’re given an integer array nums and an integer k. Your task is to count how many contiguous, non-empty subarrays have a sum exactly equal to k.

Constraints:

• $1 \leq$ nums.length $\leq 2 \times 10^4$
• $−1000 \leq$ nums[i] $\leq 1000$
• $-10^7 \leq$ k $ \leq 10^7$

Examples

Why Checking Every Subarray is Too Slow

The most direct way to solve this problem is to consider every possible subarray and compute its sum. For an array of length $n$, there are $O(n^2)$ such subarrays. Even if we optimize the sum calculation using prefix sums, we would still need to examine every $(start, end)$ pair explicitly.

With $n$ up to $20,000$, this is not feasible. So, to do better, we need a way to reason about subarray sums without enumerating all ranges explicitly.

More importantly, brute force treats each subarray as an independent object. It fails to take advantage of the fact that subarrays overlap heavily, and that much of their sum information is repeated. To solve this problem efficiently, we need a way to reuse work across overlapping subarrays instead of recomputing sums from scratch.

Optimized Approach Using Prefix Sum and Hashmap

A prefix sum represents the total sum of elements from the start of the array up to a given index. Once we view the array this way, every subarray sum can be expressed as the difference between two prefix sums.

If a subarray ending at index r has sum k, then the prefix sum just before the subarray starts must be exactly k less than the prefix sum at r. This reframes the problem: instead of searching for subarrays directly, we can search for pairs of prefix sums whose difference is k.

As we iterate through the array from left to right, we keep track of how many times each prefix sum has appeared so far. When we reach a new index, we can immediately determine how many valid subarrays end at that index by checking how often the required prefix sum has already occurred.

The hashmap is what makes this efficient. It allows us to answer the question “How many times have I seen this prefix sum before?” in constant time. Combined with prefix sums, this turns the problem into a single linear pass.

Why this works even with negatives
Unlike the sliding window (which relies on monotonic behavior), prefix sums and frequency counting using hashmap do not care whether numbers are positive, negative, or zero because it do not rely on sums increasing or decreasing in any predictable way.

Algorithm Steps

Let’s look at the algorithmic steps of the approach we just discussed.

1. Initialize a hashmap, prefixCount, to store the frequency of prefix sums.
2. Insert {0: 1} into the hashmap to represent the empty prefix before the array starts.
3. Initialize, prefixSum = 0, to track the running sum, and result = 0, to store the total number of valid subarrays
4. Iterate through each element num in nums:
   1. Add num to prefixSum
   2. If (prefixSum - k) exists in prefixCount, add its frequency to result
   3. Increment the frequency of prefixSum in the hashmap
5. Return result after processing all elements

To better understand the algorithm, let’s walk through an illustration example below:

1 / 5

Python Implementation

Let’s look at the code for the solution we just discussed.

from typing import List
from collections import defaultdict

def subarraySum(nums, k):
    # Dictionary to store how many times each prefix sum has appeared
    # prefixCount[sum] = frequency of that prefix sum
    prefixCount = defaultdict(int)
    
    # Initialize with prefix sum 0 occurring once
    # This handles subarrays that start from index 0
    prefixCount[0] = 1

    # Running prefix sum of elements processed so far
    prefixSum = 0

    # Stores the total number of subarrays whose sum equals k
    result = 0

    # Iterate through each element in the array
    for num in nums:
        # Update the running prefix sum
        prefixSum += num

        # If prefixSum - k has appeared before,
        # then there exist subarrays ending at this index with sum k
        result += prefixCount[prefixSum - k]

        # Record the current prefix sum in the hashmap
        # for future subarray calculations
        prefixCount[prefixSum] += 1

    # Return the total count of valid subarrays
    return result


def main():
    # List of test cases where each entry is (nums, k)
    test_cases = [
        ([3, 4, 7], 7),
        ([1, -1, 1], 1),
        ([0, 0, 0], 0),
        ([2, 2, 2], 1),
        ([1, 2, 3, -3, 4], 3)
    ]

    # Run the algorithm for each test case
    for idx, (nums, k) in enumerate(test_cases, start=1):
        result = subarraySum(nums, k)

        # Display formatted input and output for clarity
        print(f"{idx}.\tInput nums: [ {', '.join(str(n) for n in nums)} ], k = {k}")
        print(f"\tNumber of subarrays with sum {k} = {result}")
        print("-" * 75)


# Entry point of the program
if __name__ == "__main__":
    main()

Time Complexity

The time complexity of this solution is $O(n)$, where $n$ is the length of the input array, nums. The algorithm iterates through the array exactly once, and all hashmap lookups and updates are constant-time operations on average. As a result, the total amount of work performed grows linearly with the size of the input.

Space Complexity

The space complexity of this solution is $O(n)$. In the worst case, the hashmap may store up to $n$ distinct prefix sums if all cumulative sums are unique. So, the extra memory used by the algorithm scales linearly with the size of the input array.

Edge Cases to Consider

Let’s discuss important edge cases that should be considered while solving this problem:

• Subarrays starting at index 0
  • Example: nums = [3], k = 3 → should count [3]
  • That’s why we set prefix_count[0] = 1 at the start.
• Negative numbers
  • Example: nums = [1, -1, 1], k = 1 → multiple overlapping answers
• Many zeros
  • Example: nums = [0,0,0], k = 0 → output is 6 (all non-empty subarrays)

Mistakes to Avoid

Following are the known mistakes to avoid:

• Using sliding window: fails when negatives exist.
• Forgetting prefix_count[0] = 1: misses subarrays that start at index 0.
• Updating prefix_count[prefix_sum] before counting: can accidentally count “empty-length” or wrong ranges.