Problem Description

You are given an array of integers nums and an integer k. Your task is to find the total number of subarrays whose sum equals to k.

A subarray is a contiguous non-empty sequence of elements within an array. This means the elements must be consecutive in the original array.

For example:

• If nums = [1, 2, 3] and k = 3, the subarrays with sum equal to 3 are [1, 2] and [3], so the answer would be 2.
• The subarray must contain at least one element (non-empty).
• The elements in the subarray must appear consecutively in the original array.

The solution uses a hash table combined with prefix sum technique. The key insight is that if we have a prefix sum s at current position and we previously had a prefix sum of s - k at some earlier position, then the subarray between those two positions has a sum of exactly k.

The algorithm works as follows:

1. Initialize a counter cnt with {0: 1} to handle cases where the prefix sum itself equals k
2. Iterate through each number in the array, maintaining a running sum s
3. For each position, check how many times we've seen the prefix sum s - k before (this tells us how many subarrays ending at current position have sum k)
4. Add the current prefix sum s to our counter for future use
5. Return the total count of valid subarrays

Intuition

The brute force approach would be to check every possible subarray and count those with sum equal to k. This would require O(n²) or O(n³) time complexity, which is inefficient for large arrays.

The key observation is that we can use prefix sums to optimize this. A prefix sum at index i represents the sum of all elements from index 0 to index i. If we know the prefix sum at two different positions, we can calculate the sum of the subarray between them.

Here's the crucial insight: If the current prefix sum is s and we want a subarray with sum k, we need to find if there was a previous prefix sum equal to s - k. Why? Because if there was a prefix sum of s - k at some earlier position j, and the current prefix sum at position i is s, then the subarray from j+1 to i has sum: s - (s - k) = k.

For example, if nums = [1, 2, 3, 4] and k = 5:

• At index 2, prefix sum = 6 (1+2+3)
• We look for prefix sum = 6 - 5 = 1
• Prefix sum 1 exists at index 0
• So subarray from index 1 to 2 is [2, 3] with sum 5

To efficiently track how many times each prefix sum has occurred, we use a hash table. This allows us to instantly check if s - k exists and how many times it has appeared (each occurrence represents a different valid subarray ending at the current position).

We initialize the hash table with {0: 1} because a prefix sum of 0 exists before we start (representing an empty prefix). This handles cases where the prefix sum from the start equals k.

Learn more about Prefix Sum patterns.

Solution Approach

The solution uses a hash table combined with prefix sum technique to solve the problem efficiently in O(n) time.

Data Structure Used:

• Hash table cnt (Counter in Python) to store the frequency of each prefix sum value

Algorithm Steps:

1. Initialize the hash table: Create cnt with initial value {0: 1}. This represents that before processing any elements, we have seen prefix sum 0 exactly once. This is crucial for handling cases where a prefix sum from the beginning equals k.

2. Initialize variables:

   • ans = 0 to store the count of valid subarrays
   • s = 0 to maintain the running prefix sum

3. Iterate through the array: For each element x in nums:

   a. Update prefix sum: Add current element to running sum: s += x

   b. Count valid subarrays: Check if s - k exists in our hash table. If it does, add cnt[s - k] to our answer. This represents all subarrays ending at the current position with sum equal to k.

   c. Update hash table: Increment the count of current prefix sum: cnt[s] += 1. This records that we've seen this prefix sum value one more time, which can be used for future positions.

4. Return the result: After processing all elements, return ans which contains the total count of subarrays with sum equal to k.

Why this works:

• When we're at position i with prefix sum s, and we find that s - k appeared m times before, it means there are m different starting positions that create subarrays ending at position i with sum exactly k.
• By maintaining the frequency of all prefix sums seen so far, we can instantly determine how many valid subarrays end at each position.

Time Complexity: O(n) - single pass through the array Space Complexity: O(n) - hash table can store at most n+1 different prefix sums

Example Walkthrough

Let's walk through the algorithm with nums = [1, 2, 1, 2] and k = 3.

Initial Setup:

• cnt = {0: 1} (we've seen prefix sum 0 once)
• ans = 0 (no subarrays found yet)
• s = 0 (current prefix sum)

Step 1: Process nums[0] = 1

• Update prefix sum: s = 0 + 1 = 1
• Check for s - k = 1 - 3 = -2 in cnt: not found, so no subarrays ending here
• Update cnt: cnt = {0: 1, 1: 1}
• ans = 0

Step 2: Process nums[1] = 2

• Update prefix sum: s = 1 + 2 = 3
• Check for s - k = 3 - 3 = 0 in cnt: found with count 1
  • This means subarray [1, 2] from start has sum 3
• Add to answer: ans = 0 + 1 = 1
• Update cnt: cnt = {0: 1, 1: 1, 3: 1}

Step 3: Process nums[2] = 1

• Update prefix sum: s = 3 + 1 = 4
• Check for s - k = 4 - 3 = 1 in cnt: found with count 1
  • This means subarray [2, 1] (from index 1 to 2) has sum 3
• Add to answer: ans = 1 + 1 = 2
• Update cnt: cnt = {0: 1, 1: 1, 3: 1, 4: 1}

Step 4: Process nums[3] = 2

• Update prefix sum: s = 4 + 2 = 6
• Check for s - k = 6 - 3 = 3 in cnt: found with count 1
  • This means subarray [1, 2] (from index 2 to 3) has sum 3
• Add to answer: ans = 2 + 1 = 3
• Update cnt: cnt = {0: 1, 1: 1, 3: 1, 4: 1, 6: 1}

Result: The algorithm returns ans = 3

The three subarrays with sum 3 are:

1. [1, 2] at indices 0-1
2. [2, 1] at indices 1-2
3. [1, 2] at indices 2-3

Solution Implementation

Time and Space Complexity

The time complexity is O(n), where n is the length of the array nums. The algorithm iterates through the array exactly once with a single for loop. Within each iteration, all operations are constant time: adding to the running sum s, accessing and updating the Counter dictionary cnt[s - k] and cnt[s], and incrementing ans. Dictionary operations (lookup and insertion) in Python have an average time complexity of O(1).

The space complexity is O(n), where n is the length of the array nums. In the worst case, when all prefix sums are unique, the Counter dictionary cnt will store up to n + 1 distinct prefix sum values (including the initial 0). Each entry in the dictionary requires O(1) space, resulting in a total space complexity of O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Forgetting to Initialize the Hash Map with {0: 1}

The Problem: A common mistake is initializing an empty hash map instead of starting with {0: 1}. This causes the algorithm to miss subarrays that start from index 0 and have sum equal to k.

Example of Incorrect Code:

def subarraySum(self, nums: List[int], k: int) -> int:
    prefix_sum_count = Counter()  # Wrong! Missing {0: 1}
    result = 0
    current_sum = 0
  
    for num in nums:
        current_sum += num
        result += prefix_sum_count[current_sum - k]
        prefix_sum_count[current_sum] += 1
  
    return result

Why This Fails:

• For nums = [3, 2] and k = 5, the correct answer is 1 (subarray [3, 2])
• When processing index 1, current_sum = 5 and we check for prefix_sum_count[5 - 5] = prefix_sum_count[0]
• Without the initial {0: 1}, this returns 0 instead of 1, missing the valid subarray

Solution: Always initialize with prefix_sum_count = Counter({0: 1}) to handle subarrays starting from the beginning.

Pitfall 2: Updating Hash Map Before Checking for Valid Subarrays

The Problem: Another mistake is updating the hash map with the current prefix sum before checking for valid subarrays. This can lead to counting invalid subarrays (specifically, empty subarrays).

Example of Incorrect Code:

def subarraySum(self, nums: List[int], k: int) -> int:
    prefix_sum_count = Counter({0: 1})
    result = 0
    current_sum = 0
  
    for num in nums:
        current_sum += num
        prefix_sum_count[current_sum] += 1  # Wrong order!
        result += prefix_sum_count[current_sum - k]
  
    return result

Why This Fails:

• When k = 0, this approach incorrectly counts empty subarrays
• For nums = [0] and k = 0, after processing the element:
  • current_sum = 0
  • We increment prefix_sum_count[0] to 2
  • Then we check prefix_sum_count[0 - 0] = 2, giving us 2 instead of 1
• The extra count comes from incorrectly considering an empty subarray

Solution: Always check for valid subarrays first, then update the hash map:

result += prefix_sum_count[current_sum - k]  # Check first
prefix_sum_count[current_sum] += 1           # Then update

Pitfall 3: Using Regular Dictionary Instead of Counter

The Problem: Using a regular dictionary and not handling missing keys properly can cause KeyError exceptions.

Example of Incorrect Code:

def subarraySum(self, nums: List[int], k: int) -> int:
    prefix_sum_count = {0: 1}  # Regular dict
    result = 0
    current_sum = 0
  
    for num in nums:
        current_sum += num
        result += prefix_sum_count[current_sum - k]  # KeyError if key doesn't exist!
        if current_sum in prefix_sum_count:
            prefix_sum_count[current_sum] += 1
        else:
            prefix_sum_count[current_sum] = 1
  
    return result

Solution: Either use Counter which returns 0 for missing keys, or use dict.get() with a default value:

# Option 1: Use Counter (recommended)
prefix_sum_count = Counter({0: 1})
result += prefix_sum_count[current_sum - k]  # Returns 0 if key doesn't exist

# Option 2: Use dict.get()
prefix_sum_count = {0: 1}
result += prefix_sum_count.get(current_sum - k, 0)