Prefix Sums and Hash Maps for Subarray Sums

In our previous lesson, we established how to use a prefix sum array to answer any range sum query in constant time after an initial linear-time precomputation. This technique is a powerful building block, and today we'll see how it can be extended to solve a more dynamic and common category of interview problems.

This lesson focuses on the learning outcome: Combine prefix sums with a hash map to count subarrays with a given target sum. We will tackle the classic "Subarray Sum Equals K" problem, a staple in coding interviews that elegantly blends the prefix sum concept with the fast look-up capabilities of a hash map.

From Querying Ranges to Finding a Target Sum

In the last lesson, the problem was "given a range [L, R], find its sum." Now, we're flipping the problem on its head: "given a target sum k, find how many contiguous subarrays have this sum."

A brute-force approach would be to generate every possible subarray, calculate its sum, and check if it equals k. Generating all subarrays takes $O(N^2)$ time, and for each, summing the elements can take up to $O(N)$ time, leading to a naive $O(N^3)$ solution. We can optimize this to $O(N^2)$ by reusing the sum from the previous subarray. However, for typical interview constraints where $N$ can be up to $10^4$ or $10^5$, an $O(N^2)$ solution is still too slow. We need a linear-time approach.

The following video from take U forward provides an excellent overview of the progression from the brute-force approach to a more optimized one, which sets the stage for our final, optimal solution.

You might wonder if a sliding window approach could work. As the NeetCode video below explains, the presence of negative numbers in the array makes a standard sliding window unreliable, because expanding the window doesn't guarantee the sum increases.

The Core Insight: An Algebraic Transformation

Let's return to the foundation of our previous lesson. The sum of a subarray from index i to j (inclusive) can be expressed using prefix sums:
$\text{sum}(i,j)=\text{prefix\_sum}[j]-\text{prefix\_sum}[i-1]$

We are looking for cases where this sum equals k:
$\text{prefix\_sum}[j]-\text{prefix\_sum}[i-1]=k$

With a bit of algebraic rearrangement, we can isolate the term representing the past:
$\text{prefix\_sum}[i-1]=\text{prefix\_sum}[j]-k$

This equation is the key. It tells us that for a subarray ending at index j to have a sum of k, there must have been a previous point i-1 where the prefix sum was exactly equal to prefix_sum[j] - k.

So, the problem transforms. As we iterate through the array from left to right, calculating the current prefix sum (prefix_sum[j]), we no longer need to look back at all possible start points. We only need to ask one question: "How many times have we previously seen a prefix sum equal to current_prefix_sum - k?"

This is where the hash map comes in. A hash map provides $O(1)$ average time complexity for lookups, making it the perfect tool to track the prefix sums we've encountered so far.

The following image provides a high-level visual of this exact logic.

The Algorithm: Prefix Sums and a Hash Map

Here is the step-by-step algorithm that runs in $O(N)$ time:

1. Initialize:

   • count = 0 to store the number of valid subarrays.
   • current_sum = 0 to keep track of the running prefix sum.
   • A hash map, prefix_sums, to store the frequencies of the prefix sums encountered so far.

2. The Critical Base Case: Before starting, add an entry {0: 1} to the prefix_sums map. This represents the "empty prefix" (the sum before we've processed any elements), which has a sum of 0 and occurs once. This is crucial for correctly counting subarrays that start from index 0. For such a subarray, prefix_sum[j] = k, and our formula becomes prefix_sum[j] - k = 0. The {0: 1} entry allows us to find this match.

3. Iterate through the array: For each number num in nums:
   a. Update the running sum: current_sum += num.
   b. Find the required past sum: Calculate diff = current_sum - k.
   c. Check the hash map: If diff exists as a key in prefix_sums, it means there are one or more subarrays ending at the current position with a sum of k. Add the frequency of diff to our count: count += prefix_sums.get(diff, 0).
   d. Update the hash map: Increment the frequency of the current_sum in the prefix_sums map. prefix_sums[current_sum] = prefix_sums.get(current_sum, 0) + 1.

The order of operations is vital: you must check for diff (step 3c) before updating the map with current_sum (step 3d). If you update first, you might incorrectly count an empty subarray when k = 0.

The following detailed guide from AlgoMonster provides a superb explanation of the intuition, the algorithm, and common pitfalls.

A Detailed Walkthrough

Let's trace the algorithm with an example. Consider nums = [1, 2, 3] and k = 3.

• Initial state: count = 0, current_sum = 0, prefix_sums = {0: 1}.

• num = 1 (index 0):

  • current_sum becomes 1.
  • diff = 1 - 3 = -2. Is -2 in prefix_sums? No. count remains 0.
  • Update map: prefix_sums is now {0: 1, 1: 1}.

• num = 2 (index 1):

  • current_sum becomes 1 + 2 = 3.
  • diff = 3 - 3 = 0. Is 0 in prefix_sums? Yes, its count is 1.
  • count becomes 0 + 1 = 1. (This found the subarray [1, 2]).
  • Update map: prefix_sums is now {0: 1, 1: 1, 3: 1}.

• num = 3 (index 2):

  • current_sum becomes 3 + 3 = 6.
  • diff = 6 - 3 = 3. Is 3 in prefix_sums? Yes, its count is 1.
  • count becomes 1 + 1 = 2. (This found the subarray [3], because the prefix sum up to the previous element was 3).
  • Update map: prefix_sums is now {0: 1, 1: 1, 3: 1, 6: 1}.

• Final result: count = 2. The two subarrays are [1, 2] and [3].

The following image provides another excellent trace for a different example.

For a dynamic walkthrough, the NeetCode video provides a clear, step-by-step explanation with a great example.

Python Implementation

Here is a clean Python implementation using collections.Counter, which simplifies handling missing keys by automatically providing a default value of 0.

from collections import Counter

def subarraySum(nums: list[int], k: int) -> int:
    """
    Calculates the number of contiguous subarrays whose sum equals k.
    
    Time complexity: O(N)
    Space complexity: O(N)
    """
    count = 0
    current_sum = 0
    



    # The hash map stores {prefix_sum: frequency}
    # Initialize with a prefix sum of 0 having occurred once.
    prefix_sums = Counter({0: 1})
    
    for num in nums:



        # Update the running prefix sum
        current_sum += num
        



        # Calculate the difference we need to find in our map
        diff = current_sum - k
        



        # Add the number of times we've seen the required 'diff'
        count += prefix_sums[diff]
        



        # Increment the count for the current prefix sum
        prefix_sums[current_sum] += 1
        
    return count




# Example
nums = [1, 2, 3]
k = 3
print(f"Number of subarrays with sum {k}: {subarraySum(nums, k)}") # Output: 2

nums = [1, -1, 1, -1, 1, -1]
k = 0
print(f"Number of subarrays with sum {k}: {subarraySum(nums, k)}") # Output: 9

Conclusion

In this lesson, you've leveled up your understanding of prefix sums by combining them with a hash map to solve a non-trivial problem efficiently. This pattern is a powerful tool in your algorithmic arsenal.

Key Takeaways:

• The Problem: Counting all contiguous subarrays that sum to a target k.
• The Core Insight: The problem can be transformed algebraically: if sum(i, j) = k, then prefix_sum[i-1] = prefix_sum[j] - k.
• The Solution: Iterate through the array, maintaining a running current_sum. For each current_sum, use a hash map to instantly look up how many times you've previously seen the required current_sum - k.
• The Complexity: This approach achieves $O(N)$ time and $O(N)$ space, a significant improvement over naive solutions.
• The Critical Detail: Always initialize the hash map with a count of 1 for a prefix sum of 0 ({0: 1}) to handle subarrays that start from the beginning of the array.

We've now explored problems involving sums of subarrays. In the next lesson, we'll shift our focus slightly. Instead of finding a specific sum, we will learn how to find the subarray with the maximum possible sum using another famous and elegant linear-time algorithm: Kadane's Algorithm.