Counting Subarrays with a Given Sum

Problem: Subarray Sum Equals K

First, let's be very clear about the problem. You are given an array of integers (which can include positive numbers, negative numbers, and zeros) and an integer k. Your task is to find the total number of contiguous subarrays whose elements sum up to exactly k.

For example, if the array is [1, 2, 3] and k = 3, there are two such subarrays:

1. [1, 2] (since 1 + 2 = 3)
2. [3]

The final count is 2. The takeUforward article below provides a clear definition and more examples.

Approach 1: The Brute-Force Solution (O(n³))

The most direct way to solve this is to generate every single possible subarray, calculate the sum of each one, and check if that sum equals k.

The algorithm is as follows:

1. Use a loop with index i to select the starting point of a subarray.
2. Use a nested loop with index j to select the ending point of the subarray. This defines the subarray arr[i...j].
3. Use a third loop with index m to iterate from i to j, adding up all the elements to get the sum.
4. If sum == k, increment a counter.

This method is exhaustive and guaranteed to be correct, but as you might guess, it's very slow. The triple-nested loop structure gives it a time complexity of approximately $O(n^3)$.

The video from 'take U forward' walks through this logic.

Approach 2: A Better Solution (O(n²))

We can make a simple but significant optimization. Notice that the third loop, which recalculates the sum for every subarray, is doing a lot of repeated work. For a fixed starting point i, as we extend the end point j, we are constantly re-summing the same initial elements.

We can eliminate the third loop by maintaining a running sum.

1. Use a loop with index i to select the starting point.
2. Initialize sum = 0.
3. Use a nested loop with index j starting from i. In each iteration of this inner loop:
   a. Add arr[j] to the sum.
   b. Check if the sum equals k. If it does, increment our counter.

This approach reduces the complexity to $O(n^2)$. It's a definite improvement and a good step in an interview, but for large arrays, it will likely result in a "Time Limit Exceeded" (TLE) error.

Approach 3: The Optimal Solution with Prefix Sums and a HashMap (O(n))

To get to a linear time solution, we need a new insight. This is where the prefix sum technique comes into play.

A prefix sum at index j, let's call it prefixSum[j], is the sum of all elements from the start of the array up to j. That is, prefixSum[j] = arr[0] + arr[1] + ... + arr[j].

With prefix sums, the sum of any contiguous subarray arr[i...j] can be found in $O(1)$ time:
sum(arr[i...j]) = prefixSum[j] - prefixSum[i-1]

Our goal is to find the number of pairs (i, j) such that sum(arr[i...j]) = k. Using our new formula, this means we are looking for:
prefixSum[j] - prefixSum[i-1] = k

If we rearrange this equation, we get the key insight:
prefixSum[j] - k = prefixSum[i-1]

This changes the problem entirely. Instead of searching for subarrays, we can iterate through the array once, calculating the current prefix sum at each index j. At each step, we ask: "How many times have we previously seen a prefix sum that equals currentPrefixSum - k?"

Each time we find such a past prefix sum, it means we have found a new subarray ending at j whose sum is k.

A HashMap is the perfect tool for this. We can store the prefix sums we have seen so far and their frequencies (prefix sum -> count).

The algorithm is:

1. Initialize a HashMap to store prefix sum frequencies. A crucial first step is to put (0, 1) into the map. This handles cases where a subarray starting from index 0 sums to k. Think of it as the "empty" prefix before the array begins, which has a sum of 0 and occurs once.
2. Initialize count = 0 and currentPrefixSum = 0.
3. Iterate through the array. For each element:
   a. Add the element to currentPrefixSum.
   b. Calculate the required "past" prefix sum: remove = currentPrefixSum - k.
   c. Check the map: if remove exists as a key, add its frequency (map.get(remove)) to count.
   d. Update the map: increment the frequency of currentPrefixSum. If it's not in the map, add it with a frequency of 1.
4. Return count.

Let's trace this with arr = [1, 2, 3] and k = 3.

• Initial: map = {0: 1}, count = 0, prefixSum = 0.

• i = 0, element = 1:

  • prefixSum becomes 0 + 1 = 1.
  • remove = 1 - 3 = -2. Is -2 in the map? No.
  • Update map with prefixSum = 1. map is now {0: 1, 1: 1}.

• i = 1, element = 2:

  • prefixSum becomes 1 + 2 = 3.
  • remove = 3 - 3 = 0. Is 0 in the map? Yes, its frequency is 1.
  • count becomes 0 + 1 = 1. (We found the subarray [1, 2]).
  • Update map with prefixSum = 3. map is now {0: 1, 1: 1, 3: 1}.

• i = 2, element = 3:

  • prefixSum becomes 3 + 3 = 6.
  • remove = 6 - 3 = 3. Is 3 in the map? Yes, its frequency is 1.
  • count becomes 1 + 1 = 2. (We found the subarray [3], which is the part between prefixSum=6 and the previous prefixSum=3).
  • Update map with prefixSum = 6. map is now {0: 1, 1: 1, 3: 1, 6: 1}.

• End: Final count is 2.

The images below, from the takeUforward article, illustrate these steps.

This approach is extremely powerful because it solves the problem in a single pass. The takeUforward article provides a superb written explanation of this.

For a dynamic, step-by-step walkthrough, Striver's video is an excellent resource.

Here is the complete Java code for the optimal solution:

import java.util.HashMap;

public class Solution {
    public int subarraySum(int[] nums, int k) {
        // HashMap to store prefix sums and their frequencies
        // Key: prefix sum, Value: frequency
        HashMap<Integer, Integer> prefixSumMap = new HashMap<>();
        
        // Base case: A prefix sum of 0 exists once (the empty prefix)
        prefixSumMap.put(0, 1);
        
        int count = 0;
        int currentPrefixSum = 0;
        
        for (int num : nums) {
            // Update the current prefix sum
            currentPrefixSum += num;
            
            // Calculate the prefix sum we need to find in the map
            int removeSum = currentPrefixSum - k;
            
            // If the map contains the required 'removeSum', it means we have found
            // one or more subarrays ending at the current position with a sum of k.
            if (prefixSumMap.containsKey(removeSum)) {
                count += prefixSumMap.get(removeSum);
            }
            
            // Add the current prefix sum to the map, updating its frequency.
            prefixSumMap.put(currentPrefixSum, prefixSumMap.getOrDefault(currentPrefixSum, 0) + 1);
        }
        
        return count;
    }
}

Complexity Analysis:

• Time Complexity: $O(n)$, because we iterate through the array only once. Each HashMap operation (insertion and retrieval) takes $O(1)$ on average.
• Space Complexity: $O(n)$, as in the worst case, the HashMap might store n distinct prefix sums if all prefix sums are unique.

Conclusion

In this lesson, we explored a powerful pattern for solving subarray problems. By combining the prefix sum technique with the fast lookups of a HashMap, we were able to optimize a problem from a brute-force $O(n^3)$ or $O(n^2)$ solution down to an optimal $O(n)$ single-pass solution.

Key Takeaways:

• Problem-Solving Progression: We moved from a naive brute-force method to a better, and finally to an optimal solution, a crucial skill for interviews.
• Prefix Sums: This technique is fundamental for problems involving sums of ranges or subarrays. The formula sum(i, j) = prefixSum[j] - prefixSum[i-1] is one you should commit to memory.
• HashMap for Frequencies: The key insight was to rearrange the equation to prefixSum[j] - k = prefixSum[i-1] and use a HashMap to store frequencies of past prefix sums, enabling an efficient lookup for the removeSum.
• Edge Case Handling: The map.put(0, 1) initialization is a subtle but critical detail that correctly handles subarrays starting from the beginning of the array.

We have now concluded the "Arrays & Hashing" module. You've learned several foundational patterns, including the two-pointer technique, Kadane's algorithm, and using hash sets/maps with prefix sums for efficient lookups.

In our next lesson, we will begin the Sorting & Searching Algorithms module. We'll start by implementing the fundamental sorting algorithms—Selection Sort, Bubble Sort, and Insertion Sort—to build a strong foundation for more advanced techniques like Merge Sort and Quick Sort.