Subarray Sum Equals K Problem Solution

The Subarray Sum Equals K problem requires identifying all contiguous subarrays that sum to a target value k. This classic algorithmic challenge demonstrates the power of prefix sums and hash maps for efficient array processing.

1. Problem Restatement

Clarifying questions to ask the interviewer:

1. Can the input array contain negative numbers? (This affects whether we can use sliding window techniques)
2. What is the expected time complexity constraint? (Are we aiming for O(n) or O(n log n)?)
3. Are we allowed to modify the input array? (This affects whether we can sort or mutate)
4. What should we return if no subarray exists with sum equal to k? (Count, indices, or just existence?)
5. What is the maximum size of the input array? (This helps determine if O(n²) approach is acceptable)

Restatement in my own words: Given an integer array and a target sum k, we need to count how many contiguous subarrays have a sum equal to k. A subarray is a contiguous portion of the array, and we're looking for all such subarrays that sum to exactly k.

Problem details:

• Input: Integer array nums of length n, integer k
• Output: Number of contiguous subarrays whose sum equals k
• Constraints:
  • Array length can be up to 2×10⁴ (so O(n²) approach might be too slow)
  • Array elements can be negative, zero, or positive
  • We want to count occurrences, not return the actual subarrays

Assumptions I'm making:

• The array can contain negative numbers (this is key to the approach choice)
• We're looking for a count, not the actual subarrays themselves
• The input array can be mutated if needed (though we'll avoid this when possible)
• We're not concerned with the actual indices of subarrays, just their count

Data shape and hidden structure: The problem has a hidden prefix sum structure. If we think about cumulative sums, then any subarray from index i to j has sum equal to prefix[j] - prefix[i-1]. This suggests we should look for pairs of prefix sums that differ by exactly k, which leads to a hash map-based approach.

The key insight is that instead of checking every possible subarray (O(n²)), we can use the relationship between prefix sums to find valid subarrays in O(n) time using a hash map to store previously seen prefix sums.

2. Pattern Recognition

The problem asks us to find the number of contiguous subarrays whose sum equals a target value K. This is a classic example of the "Prefix Sum with Hash Map" pattern.

The key wording that signals this pattern:

• "Subarray sum" - indicates we're looking for contiguous elements with a specific sum property
• "Count of subarrays" - suggests we need to track occurrences rather than find specific indices
• "Target sum K" - hints at a comparison operation that can be optimized with precomputed values

The core insight is that if we know the prefix sum up to index i and we want to find subarrays ending at index i with sum K, we need to look for prefix sums that equal (current_prefix_sum - K). This transforms the problem from checking all possible subarrays O(n³) to maintaining a running count of prefix sums O(n).

This pattern fits because:

1. We have a cumulative property (prefix sum) that can be reused
2. The relationship between subarray sums and prefix sums is mathematical (subarray[i,j] = prefix[j] - prefix[i-1])
3. We can use a hash map to store and quickly retrieve previously seen prefix sums

The fundamental insight is that instead of checking every subarray explicitly, we can leverage the mathematical relationship between prefix sums to efficiently count valid subarrays.

The pattern is most commonly solved using a HashMap to store frequency counts of prefix sums, enabling O(1) lookup for the required previous prefix sum values.

3. Approach Evolution

My first instinct is to use a brute force approach where I check every possible subarray and calculate its sum, but that's O(n²) because for each starting position I would need to iterate through all ending positions. Let me think about how to optimize this.

I notice that if I maintain a running prefix sum and use a hash map to store frequency counts of each prefix sum, I can reduce this to O(n) time. The key insight is that if I want to find subarrays with sum equal to k, and I have a prefix sum at position i, then I need to find how many previous prefix sums equal (current_prefix_sum - k). This is because if prefix[i] - prefix[j] = k, then the sum of elements from index j+1 to i equals k.

The optimization works by using a HashMap to store the frequency of each prefix sum encountered so far. When I'm at any position, I check how many times (current_sum - k) has occurred before. Each occurrence represents a valid subarray ending at the current position.

Let me trace through this approach:

1. I'll iterate through the array once
2. Keep track of cumulative sum so far
3. For each position, check if (cumulative_sum - k) exists in my map
4. If it does, that many subarrays ending at current position sum to k
5. Add current cumulative sum to the map for future lookups

This approach maintains the invariant that my HashMap always contains prefix sums from previous positions, and I'm counting valid subarrays efficiently by leveraging the mathematical relationship between prefix sums.

The approach is optimal because we only make one pass through the array, and each HashMap lookup/insertion is O(1) amortized. This gives us linear time complexity with a clean mathematical foundation.

The key insight is that instead of checking all subarrays (which would require nested loops), we can transform the problem into a prefix sum counting problem where we're looking for specific differences between prefix sums.

4. Data Structure Choice

For this subarray sum problem, I'll use a HashMap to store prefix sums and their frequencies.

Primary data structure: HashMap<Integer, Integer>

• Key: prefix sum (cumulative sum from index 0 to current position)
• Value: frequency count of how many times this prefix sum has occurred

Why this data structure is needed: The HashMap enables O(1) average lookup and update operations for prefix sums, which is crucial for achieving linear time complexity. Without it, we'd need to check all previous prefix sums for each position, leading to O(n²) time complexity.

Time and space impact:

• Time: O(n) for the main loop, with O(1) average HashMap operations
• Space: O(n) in worst case where all prefix sums are unique

Helper structure: I'll also use a simple integer variable currentSum to track the running sum as I iterate through the array.

The HashMap approach is optimal here because:

1. It allows us to quickly determine how many previous prefix sums differ from current sum by exactly k
2. It handles negative numbers gracefully since we're counting frequencies rather than tracking specific subarrays
3. It's more efficient than maintaining a list of all prefix sums and checking each one

This is the key insight that transforms the problem from O(n²) to O(n) time complexity.

5. Algorithm Walkthrough

Let me trace through the algorithm with a concrete example to show how it works step-by-step.

Example: nums = [1, 2, 3, -1, 2], k = 3

Step-by-step execution:

1. Initialize: map = {0: 1}, sum = 0, count = 0

   • Start with prefix sum 0 occurring once (empty prefix)
   • sum = 0, count = 0

2. Process nums[0] = 1:

   • sum = 0 + 1 = 1
   • Check if sum - k = 1 - 3 = -2 exists in map → No
   • Add current sum to map: map = {0: 1, 1: 1}
   • count = 0

3. Process nums[1] = 2:

   • sum = 1 + 2 = 3
   • Check if sum - k = 3 - 3 = 0 exists in map → Yes (count = 1)
   • count = 0 + 1 = 1
   • Add current sum to map: map = {0: 1, 1: 1, 3: 1}

4. Process nums[2] = 3:

   • sum = 3 + 3 = 6
   • Check if sum - k = 6 - 3 = 3 exists in map → Yes (count = 1)
   • count = 1 + 1 = 2
   • Add current sum to map: map = {0: 1, 1: 1, 3: 1, 6: 1}

5. Process nums[3] = -1:

   • sum = 6 + (-1) = 5
   • Check if sum - k = 5 - 3 = 2 exists in map → No
   • Add current sum to map: map = {0: 1, 1: 1, 3: 1, 6: 1, 5: 1}
   • count = 2

6. Process nums[4] = 2:

   • sum = 5 + 2 = 7
   • Check if sum - k = 7 - 3 = 4 exists in map → No
   • Add current sum to map: map = {0: 1, 1: 1, 3: 1, 6: 1, 5: 1, 7: 1}
   • count = 2

Final result: count = 2

The two subarrays that sum to k=3 are:

• [1, 2] (indices 0-1)
• [1, 2, 3, -1] (indices 0-3)

Key insights from the walkthrough:

• At each step, we're looking for a previous prefix sum that would make the current subarray sum equal to k
• The map stores how many times each prefix sum has occurred
• When we find a matching prefix sum, it means there's a subarray ending at current position that sums to k
• The algorithm correctly handles negative numbers and maintains the invariant that we only count valid subarrays

The algorithm terminates because we process each element exactly once, and the map ensures we don't double-count subarrays.

6. Java Solution

public class SubarraySumEqualsK {
    /**
     * Find the number of contiguous subarrays that sum to k.
     * Uses prefix sum technique with HashMap for O(n) time complexity.
     * 
     * @param nums Array of integers
     * @param k Target sum
     * @return Number of subarrays that sum to k
     */
    public int subarraySum(int[] nums, int k) {
        // HashMap to store frequency of prefix sums
        Map<Integer, Integer> prefixSumCount = new HashMap<>();
        
        // Initialize with 0 sum having frequency 1 (empty prefix)
        prefixSumCount.put(0, 1);
        
        int prefixSum = 0;
        int result = 0;
        
        for (int num : nums) {
            // Update running prefix sum
            prefixSum += num;
            
            // Check if (prefixSum - k) exists in map
            // If yes, it means there are subarrays ending at current position with sum k
            int targetPrefix = prefixSum - k;
            if (prefixSumCount.containsKey(targetPrefix)) {
                result += prefixSumCount.get(targetPrefix);
            }
            
            // Update frequency of current prefix sum
            prefixSumCount.put(prefixSum, prefixSumCount.getOrDefault(prefixSum, 0) + 1);
        }
        
        return result;
    }
}

The solution uses the prefix sum technique with a HashMap to efficiently count subarrays that sum to k. The key insight is that if we have a prefix sum at position i and we want to find subarrays ending at position i with sum k, we need to look for prefix sums that equal (current_sum - k). This allows us to avoid checking all possible subarrays explicitly.

7. Complexity Analysis

Time Complexity: O(n)

The algorithm makes a single pass through the input array, performing constant-time operations at each step. For each element, we:

• Calculate a running sum (O(1))
• Check if sum - k exists in our hash map (O(1) average case)
• Update the hash map with the current sum (O(1) average case)

Since we process each element exactly once and all hash map operations are O(1) amortized, the overall time complexity is O(n).

Space Complexity: O(n)

In the worst case, our hash map will store all unique prefix sums encountered so far. This happens when all prefix sums are distinct, which can occur with arrays containing only positive numbers or when the array has a specific pattern that avoids prefix sum collisions.

The space usage is directly proportional to the number of unique prefix sums, which in the worst case equals n. Therefore, the space complexity is O(n).

Why These Bounds Are Correct

The time bound is tight because we must examine each element at least once to determine if any subarray ending at that position sums to k. The space bound is tight because we may need to store all prefix sums in our hash map.

This solution is optimal for the problem as stated. We cannot achieve better than O(n) time since we must at minimum inspect each element once, and we cannot avoid storing prefix sums in the worst case since there's no way to determine subarray sums without tracking the cumulative sums.

The solution achieves the theoretical lower bound for this problem, making it both time and space optimal.

8. Edge Cases & Pitfalls

Empty and Minimal Inputs

• Empty array: Returns 0 subarrays, which is correct since there are no elements to form a sum
• Single element array: If the element equals k, returns 1; otherwise 0
• All zeros: If k = 0, returns n(n+1)/2 (all possible subarrays), otherwise 0

Duplicate Values and Repeated States

• Duplicate prefix sums: Multiple prefix sums can have the same value, which is handled correctly by the HashMap storing counts
• Negative numbers: Can create multiple valid subarrays that sum to k, requiring careful counting

Boundary Conditions

• k = 0: Special case where we're looking for subarrays that sum to zero
• All positive numbers: May have fewer valid subarrays than expected
• All negative numbers: May have more valid subarrays than expected

Common Implementation Mistakes

• Incorrect HashMap initialization: Not handling the case where prefixSum = k (should initialize with 0 -> 1)
• Off-by-one errors: Miscounting subarrays when prefixSum - k exists in map
• Not accounting for empty subarray: The prefixSum = 0 case represents an empty subarray from beginning
• Integer overflow: When summing large numbers, though not a primary concern for this problem

Key Implementation Details

// Critical initialization: prefixSum 0 occurs once (empty prefix)
map.put(0, 1);

The most subtle edge case is handling the initial prefix sum of 0, which represents an empty subarray. Without this, we miss counting subarrays that start from index 0 and sum to k.

// Common mistake - incorrect initialization:
// map.put(0, 0); // This would cause missing count of subarrays starting at index 0

9. Alternative Approaches

There are several alternative approaches to solving the subarray sum equals K problem, each with different trade-offs:

Brute Force Approach: The most straightforward alternative is to check all possible subarrays using nested loops. For each starting index, we calculate cumulative sums until we find a subarray that sums to K. This approach has O(n³) time complexity due to nested loops and O(1) space complexity. While simple to implement, it's inefficient for large inputs and doesn't scale well.

Prefix Sum with Hash Map (Optimized): This is our chosen approach, but we could also use a HashMap to store prefix sums and their frequencies. Instead of storing only the count of each prefix sum, we could store all indices where each prefix sum occurs. This would allow us to handle cases where we need to track actual subarray positions, though it increases space complexity.

Two Pointers with Sliding Window: This approach works only when all elements are positive, which is not guaranteed in this problem. If we could ensure all numbers are positive, we could use a sliding window technique with two pointers to find subarrays that sum to K. However, this approach fails when negative numbers are present because we can't determine whether to expand or contract the window.

Binary Search on Prefix Sums: If we sort prefix sums, we could potentially use binary search to find if (current_prefix_sum - K) exists. However, this approach requires maintaining sorted order, which is inefficient as we process elements, and it doesn't preserve the original array structure needed to find actual subarrays.

The chosen approach using HashMap with prefix sums is optimal because it provides O(n) time complexity and O(n) space complexity, making it both time and space efficient for this specific problem. The other alternatives either don't work with negative numbers or are less efficient in terms of time or space complexity.

10. Test Cases

public class SubarraySumEqualsK {
    public int subarraySum(int[] nums, int k) {
        // Implementation here
    }
    
    // Test cases
    public static void main(String[] args) {
        SubarraySumEqualsK solution = new SubarraySumEqualsK();
        
        // Test case 1: Basic case
        int[] nums1 = {1, 1, 1};
        int k1 = 2;
        // Expected: 2
        // Subarrays: [1,1] at indices 0-1 and [1,1] at indices 1-2
        
        // Test case 2: Negative numbers
        int[] nums2 = {1, -1, 0, 1};
        int k2 = 0;
        // Expected: 4
        // Subarrays: [1,-1], [0], [1,-1,0], [0,1]
        
        // Test case 3: Single element
        int[] nums3 = {5};
        int k3 = 5;
        // Expected: 1
        // Subarray: [5]
        
        // Test case 4: Empty array
        int[] nums4 = {};
        int k4 = 0;
        // Expected: 0
        
        // Test case 5: No valid subarrays
        int[] nums5 = {1, 2, 3};
        int k5 = 7;
        // Expected: 0
        
        // Test case 6: All zeros
        int[] nums6 = {0, 0, 0};
        int k6 = 0;
        // Expected: 6
        // All possible subarrays sum to 0
        
        // Test case 7: Large negative k
        int[] nums7 = {1, 2, 3};
        int k7 = -6;
        // Expected: 0
    }
}

Sample Test Cases:

1. Basic case: nums = [1,1,1], k = 2 → Expected: 2

   • Two subarrays: [1,1] at indices 0-1 and [1,1] at indices 1-2

2. Negative numbers: nums = [1,-1,0,1], k = 0 → Expected: 4

   • Four subarrays that sum to zero

3. Single element: nums = [5], k = 5 → Expected: 1

   • One subarray that sums to target

4. Edge case - empty array: nums = [], k = 0 → Expected: 0

5. No valid subarrays: nums = [1,2,3], k = 7 → Expected: 0

6. All zeros: nums = [0,0,0], k = 0 → Expected: 6

   • All possible subarrays sum to zero

7. Large negative k: nums = [1,2,3], k = -6 → Expected: 0

   • No subarrays sum to negative target

These test cases cover various scenarios including positive/negative numbers, edge cases with empty arrays, and boundary conditions that validate the correctness of the solution.

11. Final Summary

The core insight is to use a prefix sum approach with a HashMap to count subarrays efficiently. By maintaining a running sum and tracking how many times each prefix sum has occurred, we can determine in O(1) time whether a valid subarray ending at the current position exists. This transforms the problem from checking all possible subarrays (O(n²)) into a single pass through the array with O(n) time and space complexity.

The solution works because if we want a subarray with sum K, and our current prefix sum is currentSum, then we need to find how many previous prefix sums equal currentSum - K. This is exactly what the HashMap stores. The key optimization is recognizing that we don't need to check all subarrays explicitly, but rather leverage the mathematical relationship between prefix sums.

This approach is optimal for the given constraints since we must examine each element at least once to determine valid subarrays, making O(n) time complexity unavoidable. The space usage of O(n) is also necessary to store the frequency map of prefix sums, which is the minimal space required for this approach.