Count Subarray Sum Equals K | Brute to Optimal Strategy

The Count Subarray Sum Equals K problem requires finding the total number of contiguous subarrays within an array whose elements sum up to a target value $K$. While intuitive brute-force approaches iterate through all possible contiguous segments, optimal solutions leverage the mathematical relationship between cumulative prefix sums and target differences using a Hash Map to reduce the time complexity to linear time.

1. Problem Definition and Subarray vs. Subsequence Mechanics

To solve this problem correctly, it is essential to distinguish between a subarray and a subsequence:

• Subarray: A contiguous sequence of elements within an array. Elements must appear adjacent to one another without skipping any indices.
• Subsequence: A sequence derived by deleting zero or more elements without changing the order of the remaining elements. Elements do not need to be contiguous.

Core Requirement: Given an array of $N$ integers (which can include positive, negative, and zero values) and a target sum $K$, count the total number of contiguous subarrays whose elements sum to $K$.

2. Example Analysis and Boundary Exploration

Consider the following input array and target value:

Evaluating contiguous segments yielding a sum of $3$:

1. $[1, 2]$ (indices $0 \to 1$) $\implies 1 + 2 = 3$
2. $[1, 2, 3, -3]$ (indices $0 \to 3$) $\implies 1 + 2 + 3 + (-3) = 3$
3. $[3]$ (index $2$) $\implies 3 = 3$
4. $[2, 3, -3, 1]$ (indices $1 \to 4$) $\implies 2 + 3 - 3 + 1 = 3$
5. $[3, -3, 1, 1, 1]$ (indices $2 \to 6$) $\implies 3 - 3 + 1 + 1 + 1 = 3$
6. $[1, 1, 1]$ (indices $4 \to 6$) $\implies 1 + 1 + 1 = 3$
7. $[4, 2, -3]$ (indices $7 \to 9$) $\implies 4 + 2 - 3 = 3$
8. $[-3, 1, 1, 1, 4, 2, -3]$ (indices $3 \to 9$) $\implies -3 + 1 + 1 + 1 + 4 + 2 - 3 = 3$

Total Count = 8.

3. Naive Brute-Force Approach ($O(N^3)$)

Intuition

The baseline approach generates every possible pair of starting index $i$ and ending index $j$. For each defined range $[i \dots j]$, a third loop iterates from $k = i$ to $j$ to sum up the elements and check if the sum equals $K$.

Execution Flow

C++ Code Implementation

int countSubarraySumBrute(const std::vector<int>& arr, int K) {
    int n = arr.size();
    int count = 0;

    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            int sum = 0;
            for (int k = i; k <= j; k++) {
                sum += arr[k];
            }
            if (sum == K) {
                count++;
            }
        }
    }
    return count;
}

• Time Complexity: $O(N^3)$ — Generating all subsegment pairs takes $O(N^2)$, and calculating the sum takes $O(N)$ for each pair.
• Space Complexity: $O(1)$ — Uses only scalar counter variables.

4. Optimized Two-Loop Approach ($O(N^2)$)

Intuition

Notice that recomputing the sum from scratch for every $[i \dots j]$ range using the third loop is redundant. When $j$ moves to $j + 1$, the sum of subarray $[i \dots j+1]$ is simply the existing sum of $[i \dots j]$ plus arr[j + 1].

By incrementally building the sum within the second loop, the inner $k$ loop can be completely eliminated.

C++ Code Implementation

int countSubarraySumBetter(const std::vector<int>& arr, int K) {
    int n = arr.size();
    int count = 0;

    for (int i = 0; i < n; i++) {
        int sum = 0;
        for (int j = i; j < n; j++) {
            sum += arr[j]; // Incrementally add current element
            if (sum == K) {
                count++;
            }
        }
    }
    return count;
}

• Time Complexity: $O(N^2)$ — Two nested loops iterating over all array indices.
• Space Complexity: $O(1)$ — Auxiliary variables only.

5. Optimal Mental Model: Reverse Engineering via Prefix Sums

To achieve linear time $O(N)$, we leverage the Prefix Sum concept combined with Hash Maps.

Conceptual Break Down

Let $S$ be the total cumulative prefix sum of elements from index $0$ up to the current index $i$.

If there exists a valid contiguous subarray ending at index $i$ that sums to $K$, then the prefix sum preceding this subarray must equal $S - K$.

Mathematical Identity

Instead of searching forward for subarrays that sum to $K$, we look backward in time:

• At index $i$, calculate the current running cumulative sum $S$.
• Check how many times the prefix sum $(S - K)$ has occurred in the past.
• Every prior occurrence of $(S - K)$ represents a starting point for a valid subarray ending at $i$ with a sum equal to $K$.

6. The $S - K$ Equivalence & Hash Map Role

Because an array can contain zero or negative numbers, a specific prefix sum value can appear multiple times. Therefore, a simple boolean lookup or hash set is insufficient; we must track the frequency of every prefix sum using a Hash Map (Map<PrefixSum, Frequency>).

Workflow Diagram

7. Initializing the Edge Case: Map Entry {0: 1}

A critical requirement in prefix sum algorithms is pre-populating the Hash Map with the entry {0: 1} before processing any elements.

Why is {0: 1} required?

Suppose a subarray starts directly at index $0$ and has a sum equal to $K$.

1. At this index, the cumulative prefix sum $S$ equals $K$.
2. The search target is $S - K = K - K = 0$.
3. If $0$ is not stored in the map with a frequency of $1$, this prefix segment starting from index $0$ will be missed.

Rule: Storing {0: 1} represents a hypothetical "empty prefix" before index $0$, which has a cumulative sum of $0$ and occurs exactly once.

8. Step-by-Step Execution Log (Dry Run)

Input Data

Final Count Output = 8

9. Code Implementations (C++, Java, Python)

C++ Implementation

#include <iostream>
#include <vector>
#include <unordered_map>

int subarraySum(const std::vector<int>& nums, int k) {
    std::unordered_map<int, int> prefixMap;
    prefixMap[0] = 1; // Base case: prefix sum of 0 appears once
    
    int prefixSum = 0;
    int count = 0;

    for (int num : nums) {
        prefixSum += num;
        
        int remove = prefixSum - k;
        if (prefixMap.find(remove) != prefixMap.end()) {
            count += prefixMap[remove];
        }
        
        prefixMap[prefixSum]++;
    }

    return count;
}

Java Implementation

import java.util.HashMap;
import java.util.Map;

public class Solution {
    public int subarraySum(int[] nums, int k) {
        Map<Integer, Integer> prefixMap = new HashMap<>();
        prefixMap.put(0, 1); // Base case initialization
        
        int prefixSum = 0;
        int count = 0;

        for (int num : nums) {
            prefixSum += num;
            
            int remove = prefixSum - k;
            if (prefixMap.containsKey(remove)) {
                count += prefixMap.get(remove);
            }
            
            prefixMap.put(prefixSum, prefixMap.getOrDefault(prefixSum, 0) + 1);
        }

        return count;
    }
}

Python Implementation

def subarray_sum(nums: list[int], k: int) -> int:
    prefix_map = {0: 1}  # Base case initialization
    prefix_sum = 0
    count = 0

    for num in nums:
        prefix_sum += num
        
        remove = prefix_sum - k
        if remove in prefix_map:
            count += prefix_map[remove]
            
        prefix_map[prefix_sum] = prefix_map.get(prefix_sum, 0) + 1

    return count

10. Algorithmic Complexity & Map Trade-offs

Time Complexity

• Best/Average Case: $O(N)$ when using hash map lookup (std::unordered_map in C++, HashMap in Java, dict in Python) operating in $O(1)$ constant expected time per insertion/lookup.
• Worst Case (Hash Collisions in C++): $O(N^2)$ if severe hash collisions occur in std::unordered_map. Alternatively, using std::map (Red-Black Tree implementation) guarantees $O(N \log N)$ worst-case performance.

Space Complexity

• Worst Case: $O(N)$ auxiliary space required to store up to $N + 1$ unique prefix sums in the map data structure.

11. Strategy Comparison Matrix

12. Summary and Key Takeaways

• Subarray Definition: Must be contiguous. Non-contiguous selections are subsequences and invalid for this problem.
• Prefix Sum Strategy: Instead of calculating subsegment sums forward, compute cumulative sums $S$ and search backward for prior occurrences of prefix sums equal to $S - K$.
• Mathematical Foundation: $S_{\text{current}} - S_{\text{previous}} = K \implies S_{\text{previous}} = S_{\text{current}} - K$.
• Frequency Map: Negative numbers and zeroes allow prefix sums to repeat. A Hash Map tracking {PrefixSum: Frequency} ensures all starting positions are accounted for.
• Crucial Base Case: Initialize the hash map with {0: 1} to handle scenarios where a prefix sum itself matches $K$ directly from index $0$.
• Efficiency Standard: Achieves $O(N)$ linear time complexity and $O(N)$ auxiliary space complexity, providing the optimal solution for competitive programming and interview standards.