Subarray Sum Equals K Problem Solution

The subarray sum equals k problem requires identifying all contiguous subarrays that sum to a target value. This classic algorithmic challenge demonstrates the power of prefix sum techniques combined with hash map optimizations for efficient lookup and counting.

1. Pattern Restatement

Clarifying questions to ask the interviewer:

1. What are the constraints on input array size? (e.g., up to 10⁵ elements)
2. Can the input array contain negative numbers? (Yes, this is key to the problem)
3. Should we return the count of subarrays or the actual subarrays themselves?
4. Are we allowed to modify the input array in-place?
5. What is the expected time complexity for follow-up? (O(n) preferred)

Restate the pattern or problem type in your own words: The Subarray Sum Equals K problem asks us to find the number of contiguous subarrays in an array that sum to a given target value K. This is fundamentally about identifying ranges where cumulative sums have a specific difference relationship.

Clarify the signal words and constraint shape that point to it:

• "Subarray" or "contiguous elements"
• "Sum equals K" or "target sum"
• Array input with potentially negative numbers
• Need to count occurrences rather than return actual subarrays

Note any assumptions you are making:

• The input array can contain negative numbers (this makes it more complex than just positive numbers)
• We're counting occurrences, not returning actual subarrays
• Array indices are 0-based and valid

Identify whether the pattern is best suited for static, streaming, or online scenarios: This pattern is primarily suited for static scenarios where the entire input array is available upfront. The sliding window approach works well for static arrays with non-negative numbers, but this problem requires handling negative numbers, making it better suited for the prefix sum approach that works with static data. It's not naturally suited for streaming or online processing where elements arrive one at a time.

The key insight is that if we know prefix sums at two positions, their difference gives us the sum of elements between those positions. This transforms the problem from finding subarrays with specific sums to finding pairs of prefix sums that differ by K.

2. When to Use This Pattern

Signal words and constraint shapes:

• "Subarray" or "contiguous elements"
• "Sum equals K" or "target sum"
• Array of integers (often with negatives)
• Need to find count of subarrays with specific sum property
• "Number of subarrays" or "ways to partition" often indicate this pattern

Why this pattern fits better than alternatives: This pattern is specifically suited for problems where we need to count or find contiguous subarrays with a target sum. Unlike sliding window (which requires positive numbers) or two pointers (which don't naturally handle sum counting), this approach uses prefix sums with hash maps to achieve linear time complexity. The key insight is that if we know the prefix sum up to index i and want subarrays ending at i with sum K, we can look for prefix sums that equal (current_prefix_sum - K).

When variations apply:

• Maximum subarray sum: Use Kadane's algorithm instead
• Subarray with sum closest to K: Use sliding window or two pointers with binary search
• Subarray with sum divisible by K: Use modular arithmetic with prefix sums
• Multiple target sums: May require extension to handle multiple queries

Core insight that makes this pattern work: The key insight is that for any subarray from index j+1 to i, its sum equals (prefix_sum[i] - prefix_sum[j]). If we want this sum to equal K, then we need prefix_sum[j] = (prefix_sum[i] - K). This allows us to transform the problem from finding subarrays with sum K into counting occurrences of specific prefix sums, which can be efficiently tracked using a hash map.

This pattern is particularly powerful because it transforms a quadratic problem (checking all subarrays) into a linear one by leveraging the mathematical relationship between prefix sums and subarray sums.

3. Approach Evolution

Interview communication script:

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

I notice that if I maintain a running prefix sum, I can reframe the problem. Instead of checking all subarrays, I can think: for any subarray from index i to j, its sum equals (prefixSum[j] - prefixSum[i-1]). If this sum should equal k, then I need prefixSum[j] - prefixSum[i-1] = k, which means prefixSum[i-1] = prefixSum[j] - k.

This insight suggests that if I keep track of how many times each prefix sum has occurred, I can quickly determine how many subarrays ending at the current position have sum k. This is essentially a two-sum variant where I'm looking for previous prefix sums that differ by exactly k.

Let me think about the invariants that must hold:

1. The prefix sum array must be maintained correctly as we iterate
2. For each position, I need to look back at all previously seen prefix sums that satisfy the condition
3. The count of valid subarrays should only increase when we find matching prefix sums

Brute-force approach and why it is not ideal:

The brute force approach would involve checking every possible subarray by computing its sum. For an array of size n, we have n(n+1)/2 subarrays to check. For each subarray from index i to j, we compute the sum of elements from i to j. This results in O(n³) time complexity when we don't optimize the sum calculation, or O(n²) if we use prefix sums but still check all pairs.

Improved approach and the key optimization:

The key optimization is to use a hash map to store prefix sums. As we iterate through the array, we maintain a running sum and check if (current_sum - k) has been seen before. If so, it means there are subarrays ending at the current position with sum k.

The crucial insight is that instead of checking all possible subarrays, we can compute how many valid starting points exist for subarrays ending at the current position by looking up previous prefix sums that differ by exactly k.

Final approach with step-by-step reasoning:

1. Initialize a hash map to store frequency of prefix sums, with initial entry {0: 1} (empty prefix)
2. Iterate through the array while maintaining a running sum
3. For each element, check if (current_sum - k) exists in the map
4. If it exists, add its frequency to our result count
5. Update the map with current prefix sum frequency
6. Continue until all elements are processed

Important invariants that must hold:

1. The hash map always contains the frequency of all prefix sums seen so far
2. For any position i, we've already processed all positions before i
3. When we look up (current_sum - k), we're guaranteed to find only valid previous prefix sums that could form subarrays ending at current position
4. The initial entry {0: 1} handles the case where a subarray starting from index 0 has sum k

4. Data Structure Choice

For the subarray sum equals K 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 needed: This allows O(1) lookup to check if (current_sum - K) exists in previous prefix sums
• Time impact: O(1) average lookup and insertion
• Space impact: O(n) in worst case where all prefix sums are unique

Helper Data Structure: Integer variable

• Current running sum from array start
• Used to calculate prefix sums incrementally
• Constant space overhead

The HashMap is essential because it transforms the problem from checking all possible subarrays (O(n²)) to just checking if a complement exists in previously seen prefix sums. This is the core optimization that makes the solution linear time.

Alternative: Using a HashSet would only store prefix sums without frequencies, but we need to track how many times each sum appears to handle cases with multiple valid subarrays. A TreeSet would provide ordered prefix sums but adds unnecessary complexity for this problem's requirements.

The HashMap approach is optimal here because it supports both fast lookups and frequency tracking in a single data structure, enabling the linear-time solution that's the hallmark of this pattern.

5. Algorithm Walkthrough

Let me trace through the algorithm using the example nums = [1, 2, 3, 1, 2, 3] with k = 3.

Step 1: Initialize prefixSum = 0, count = 0, and map = {0:1} (empty prefix).

Step 2: Process nums[0] = 1:

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

Step 3: Process nums[1] = 2:

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

Step 4: Process nums[2] = 3:

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

Step 5: Process nums[3] = 1:

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

Step 6: Process nums[4] = 2:

• prefixSum = 7 + 2 = 9
• Check if prefixSum - k = 9 - 3 = 6 exists in map → Yes, count = 1
• count = 2 + 1 = 3
• Add prefixSum = 9 to map: map = {0:1, 1:1, 3:1, 6:1, 7:1, 9:1}

Step 7: Process nums[5] = 3:

• prefixSum = 9 + 3 = 12
• Check if prefixSum - k = 12 - 3 = 9 exists in map → Yes, count = 1
• count = 3 + 1 = 4
• Add prefixSum = 12 to map: map = {0:1, 1:1, 3:1, 6:1, 7:1, 9:1, 12:1}

Final result is count = 4.

The algorithm correctly identifies that there are 4 subarrays with sum equal to k:

• [1,2] at indices 0-1
• [2,3,1] at indices 1-3
• [3,1,2] at indices 2-4
• [1,2,3] at indices 3-5

The key insight is that for any current prefix sum, if we've seen a previous prefix sum such that their difference equals k, then the subarray between those two points has sum k. The map efficiently tracks all previous prefix sums to make this check in O(1) time.

graph TD
    A[Start] --> B[Initialize prefixSum=0, count=0, map={0:1}]
    B --> C[Process nums[0]=1]
    C --> D[prefixSum=1]
    D --> E[Check prefixSum-k=-2 in map? No]
    E --> F[Add prefixSum=1 to map]
    F --> G[Count=0]
    G --> H[Process nums[1]=2]
    H --> I[prefixSum=3]
    I --> J[Check prefixSum-k=0 in map? Yes, count=1]
    J --> K[Count=1]
    K --> L[Add prefixSum=3 to map]
    L --> M[Process nums[2]=3]
    M --> N[prefixSum=6]
    N --> O[Check prefixSum-k=3 in map? Yes, count=1]
    O --> P[Count=2]
    P --> Q[Add prefixSum=6 to map]
    Q --> R[Process nums[3]=1]
    R --> S[prefixSum=7]
    S --> T[Check prefixSum-k=4 in map? No]
    T --> U[Count=2]
    U --> V[Add prefixSum=7 to map]
    V --> W[Process nums[4]=2]
    W --> X[prefixSum=9]
    X --> Y[Check prefixSum-k=6 in map? Yes, count=1]
    Y --> Z[Count=3]
    Z --> AA[Add prefixSum=9 to map]
    AA --> AB[Process nums[5]=3]
    AB --> AC[prefixSum=12]
    AC --> AD[Check prefixSum-k=9 in map? Yes, count=1]
    AD --> AE[Count=4]
    AE --> AF[Add prefixSum=12 to map]
    AF --> AG[Return count=4]

6. Java Solution

public class SubarraySumEqualsK {
    /**
     * Find the number of 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 sum 0 having frequency 1 (empty prefix)
        prefixSumCount.put(0, 1);
        
        int prefixSum = 0;
        int result = 0;
        
        // Iterate through each element
        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;
    }
}

This solution implements the prefix sum approach with HashMap optimization. The key insight is that for any subarray ending at current position with sum k, there must exist a previous prefix sum such that (current_prefix_sum - previous_prefix_sum) = k, which means previous_prefix_sum = (current_prefix_sum - k).

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 prefix sum (O(1))
• Check if currentSum - k exists in our HashMap (O(1) average case)
• Update the HashMap with the current prefix sum (O(1) average case)

Since each operation inside the loop is O(1) and we perform n such operations, the overall time complexity is O(n).

Space Complexity: O(n)

In the worst case, our HashMap will store all possible prefix sums we've encountered so far. This happens when all prefix sums are distinct, which can occur in arrays with all positive numbers or when there are no subarrays that sum to k.

The space usage is dominated by the HashMap storing at most n prefix sums, making the space complexity O(n).

Why These Bounds Are Correct

The bounds are tight because:

1. We must examine each element at least once to determine all valid subarrays
2. The HashMap operations are O(1) amortized due to the hash table's load factor and collision handling
3. The maximum number of distinct prefix sums we can encounter is n

Optimality

This solution is optimal for the problem as stated. We cannot achieve better than O(n) time complexity since we must inspect each element at least once to identify all subarrays that sum to k. The space complexity is also optimal since we must store at least the prefix sums we've seen to efficiently compute the answer.

The solution represents a perfect balance between time and space efficiency for this canonical subarray sum problem.

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: Only checks if that element equals k, which is valid for one subarray
• Array with zero elements: Same as empty array case

Duplicate Values and Repeated States

• Duplicate subarray sums: Multiple subarrays can sum to the same value k, so we must count all occurrences
• Zero elements in array: Zeros can be part of multiple valid subarrays, and we must not skip them
• Negative numbers: Negative values can create multiple valid subarrays that sum to k, and we must handle them correctly

Boundary Conditions

• k = 0: Need to handle when subarrays sum to zero, including cases where prefix sums are equal
• All positive numbers: Should still work correctly with proper prefix sum tracking
• All negative numbers: Valid edge case that requires careful handling of prefix sums

Common Implementation Mistakes

• Incorrect prefix sum calculation: Using wrong indices when calculating prefix sums, leading to wrong subarray sums
• Not handling zero prefix sum: Failing to account for subarrays starting from index 0 that sum to k
• Off-by-one errors: Miscounting subarrays due to incorrect bounds in prefix sum comparisons
• Using wrong data structure: Using a regular HashMap instead of a HashMap with default value handling for missing keys

Key Implementation Notes

• The algorithm relies on the mathematical property that sum[i..j] = prefixSum[j] - prefixSum[i-1]
• When prefixSum[j] - prefixSum[i-1] = k, we have found a valid subarray
• We must handle the case where prefixSum[j] = k (subarray from index 0 to j)
• The HashMap should handle missing keys gracefully by defaulting to 0 count

// Example edge case handling:
// Input: nums = [1, 2, 3], k = 3
// Expected: 2 subarrays ([1,2] and [3])
// Edge case: nums = [0, 0, 0], k = 0
// Expected: 6 subarrays (all combinations of zeros)

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, compute cumulative sums until we find a subarray that sums to K. This approach has O(n³) time complexity due to computing cumulative sums for each subarray, and O(1) space complexity. While simple to implement, it's inefficient for large inputs.

Two Pointers with Sliding Window: This approach works well when all numbers are positive, allowing us to use a sliding window technique. We maintain a window and expand/contract it based on whether the current sum is less than or greater than K. However, this approach fails when negative numbers are present because we cannot guarantee that expanding the window will increase the sum or that contracting it will decrease the sum.

Prefix Sum with HashMap (Optimal): Our chosen solution uses prefix sums with a HashMap to store cumulative sums, achieving O(n) time and space complexity. This is optimal for the general case with negative numbers.

Binary Search on Prefix Sums: If we could sort prefix sums, we might use binary search to find valid subarrays. However, this approach requires maintaining sorted prefix sums, which adds complexity and doesn't improve the overall time complexity.

The key insight is that the HashMap-based approach is optimal because it leverages the mathematical property that if prefix[i] - prefix[j] = K, then prefix[i] = prefix[j] + K. This allows us to count valid subarrays in a single pass without needing to compute all subarray sums explicitly.

The HashMap approach is preferred over alternatives because it maintains linear time complexity while handling all edge cases including negative numbers, duplicates, and zero values. Other approaches either have worse time complexity or fail to handle the general case properly.

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] and [1,1] (at indices 0-1 and 1-2)
        
        // Test case 2: Negative numbers
        int[] nums2 = {1, -1, 0, 1};
        int k2 = 0;
        // Expected: 4
        // Subarrays [1,-1], [0], [1], [1,-1,0]
        
        // Test case 3: Empty array
        int[] nums3 = {};
        int k3 = 0;
        // Expected: 0
        
        // Test case 4: Single element matching k
        int[] nums4 = {5};
        int k4 = 5;
        // Expected: 1
        
        // Test case 5: No matching 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 with sum 2: [1,1] at indices (0,1) and (1,2)

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

   • Four subarrays with sum 0: [1,-1], [0], [1], [1,-1,0]

3. Empty array: nums = [], k = 0 → Expected: 0

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

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

Tricky Test Case:

• All zeros: nums = [0,0,0], k = 0 → Expected: 6
  • All possible subarrays sum to zero, giving us 3+2+1 = 6 combinations

Boundary Test Case:

• Large negative k: nums = [1,2,3], k = -6 → Expected: 0
  • No subarray can sum to negative value when all elements are positive

11. Final Summary

The core idea is to leverage a prefix sum approach with a hash map to efficiently count subarrays that sum to a target value K. By maintaining a running cumulative sum and tracking how many times each prefix sum has occurred, we can determine in O(1) time whether a valid subarray exists ending at the current position. This works because if currentSum - targetSum has been seen before, then there are subarrays ending at the current position with sum K.

The solution is optimal with O(n) time and O(n) space complexity, making it ideal for large input sizes. The key insight is that we're not just looking for a single subarray, but counting all possible subarrays that meet the criteria, which makes the hash map approach superior to brute force methods.

What makes this solution efficient is the transformation of a nested loop problem into a single pass with constant-time lookups, turning an O(n²) brute force approach into an optimal O(n) solution. The hash map serves as a lookup table that remembers all previous prefix sums, enabling us to quickly identify valid subarray combinations without recalculating sums.