Problem Description

You are given a binary array nums that contains only 0s and 1s. Your task is to find the maximum length of a contiguous subarray that contains an equal number of 0s and 1s.

For example, if nums = [0, 1, 0, 0, 1, 1, 0], the longest contiguous subarray with equal 0s and 1s would be [0, 1, 0, 0, 1, 1] with length 6 (containing three 0s and three 1s).

The key insight is to transform this problem using prefix sums. By treating each 0 as -1 and each 1 as +1, we can calculate a running sum as we traverse the array. When this running sum reaches the same value at two different positions, it means the subarray between those positions has an equal number of 0s and 1s (since the net change is zero).

The solution uses a hash table to track where each prefix sum value first appears. As we iterate through the array:

• We update the running sum: add 1 for each 1, subtract 1 for each 0
• If we've seen this sum before, we've found a balanced subarray from the previous occurrence to the current position
• We keep track of the maximum length found

The hash table is initialized with {0: -1} to handle the case where a balanced subarray starts from index 0.

Intuition

The challenge is to find subarrays where the count of 0s equals the count of 1s. A direct approach would be to check every possible subarray and count 0s and 1s, but this would be inefficient.

Let's think about what it means for a subarray to have equal 0s and 1s. If we have equal counts, then count(1) - count(0) = 0. This gives us a key insight: instead of tracking two separate counts, we can track a single value that represents the difference.

What if we transform the problem? We can treat each 0 as -1 and keep 1 as +1. Now, a subarray with equal 0s and 1s would have a sum of 0. For example, [0, 1, 0, 1] becomes [-1, 1, -1, 1], which sums to 0.

This transformation leads us to prefix sums. As we traverse the array, we maintain a running sum s. When we add +1 for a 1 and -1 for a 0, something interesting happens: if the prefix sum at two different positions is the same, the subarray between them must sum to 0 (meaning equal 0s and 1s).

Why? Consider positions i and j where j < i. If prefix_sum[i] = prefix_sum[j], then the sum of elements from j+1 to i is prefix_sum[i] - prefix_sum[j] = 0. This means the subarray from j+1 to i has equal 0s and 1s.

To efficiently find when we've seen a prefix sum before, we use a hash table to store each prefix sum and its first occurrence index. When we encounter a prefix sum we've seen before, we can immediately calculate the length of the balanced subarray. We initialize the hash table with {0: -1} to handle cases where the balanced subarray starts from the beginning of the array.

Learn more about Prefix Sum patterns.

Solution Approach

The implementation uses the prefix sum technique combined with a hash table to efficiently find the maximum length subarray with equal 0s and 1s.

Step-by-step Implementation:

1. Initialize the hash table: We create a dictionary d = {0: -1} that maps prefix sum values to their first occurrence indices. The initial entry {0: -1} handles the case where a valid subarray starts from index 0.

2. Initialize variables: Set ans = 0 to track the maximum length found, and s = 0 as our running prefix sum.

3. Iterate through the array: For each element at index i:

   • Update the prefix sum: If the current element x is 1, add 1 to s. If it's 0, subtract 1 from s (treating 0 as -1).

   s += 1 if x else -1

4. Check if prefix sum exists in hash table:

   • If s is already in the hash table: This means we've seen this prefix sum before at index d[s]. The subarray from d[s] + 1 to i has a sum of 0, meaning equal 0s and 1s. Calculate its length as i - d[s] and update ans if this length is greater.

   if s in d:
       ans = max(ans, i - d[s])

   • If s is not in the hash table: Store this prefix sum and its current index for future reference.

   else:
       d[s] = i

5. Return the result: After processing all elements, ans contains the maximum length of a contiguous subarray with equal 0s and 1s.

Why this works:

• When the prefix sum at position i equals the prefix sum at an earlier position j, the subarray between them (from j+1 to i) must have a net sum of 0.
• By storing only the first occurrence of each prefix sum, we maximize the potential length of subarrays.
• The time complexity is O(n) as we traverse the array once, and space complexity is O(n) for the hash table.

Example Walkthrough

Let's trace through a small example to illustrate the solution approach.

Example: nums = [0, 1, 0, 1]

Initial Setup:

• Hash table d = {0: -1} (maps prefix sum to first occurrence index)
• Maximum length ans = 0
• Running sum s = 0

Step-by-step Execution:

Index 0: nums[0] = 0

• Transform 0 to -1, so s = 0 + (-1) = -1
• Check if s = -1 exists in hash table: No
• Add to hash table: d = {0: -1, -1: 0}
• Current ans = 0

Index 1: nums[1] = 1

• Add 1 to sum: s = -1 + 1 = 0
• Check if s = 0 exists in hash table: Yes! (at index -1)
• Calculate subarray length: 1 - (-1) = 2
  • This represents subarray from index 0 to 1: [0, 1]
• Update ans = max(0, 2) = 2

Index 2: nums[2] = 0

• Transform 0 to -1, so s = 0 + (-1) = -1
• Check if s = -1 exists in hash table: Yes! (at index 0)
• Calculate subarray length: 2 - 0 = 2
  • This represents subarray from index 1 to 2: [1, 0]
• Keep ans = max(2, 2) = 2

Index 3: nums[3] = 1

• Add 1 to sum: s = -1 + 1 = 0
• Check if s = 0 exists in hash table: Yes! (at index -1)
• Calculate subarray length: 3 - (-1) = 4
  • This represents the entire array from index 0 to 3: [0, 1, 0, 1]
• Update ans = max(2, 4) = 4

Result: The maximum length is 4, which corresponds to the entire array [0, 1, 0, 1] containing two 0s and two 1s.

Key Observations:

• When we encountered prefix sum 0 at index 1, it meant the subarray [0, 1] was balanced
• When we saw prefix sum 0 again at index 3, it meant the entire array from start was balanced
• The hash table stores only the first occurrence of each prefix sum to maximize potential subarray lengths

Solution Implementation

Time and Space Complexity

The time complexity is O(n), where n is the length of the input array nums. The algorithm iterates through the array exactly once using the enumerate function, and for each element, it performs constant-time operations: updating the sum s, checking if s exists in the dictionary d, calculating the maximum, and potentially adding an entry to the dictionary. Dictionary operations (lookup and insertion) are O(1) on average.

The space complexity is O(n), where n is the length of the input array. In the worst case, when all prefix sums are unique, the dictionary d will store up to n different key-value pairs (one for each possible prefix sum value). Each entry in the dictionary takes constant space, resulting in O(n) total space usage. The remaining variables (ans, s, i, x) use only O(1) additional space.

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

Common Pitfalls

1. Overwriting Earlier Prefix Sum Indices

A critical mistake is updating the hash table every time we encounter a prefix sum, rather than only storing its first occurrence:

Incorrect Implementation:

for index, value in enumerate(nums):
    cumulative_sum += 1 if value == 1 else -1
    if cumulative_sum in sum_to_index:
        current_length = index - sum_to_index[cumulative_sum]
        max_length = max(max_length, current_length)
    # WRONG: This overwrites the earlier index!
    sum_to_index[cumulative_sum] = index  # Should be inside else block

Why this fails: By constantly updating the index, we lose the earliest occurrence of each prefix sum, resulting in shorter subarray lengths. For example, with nums = [0,1,0,1], the prefix sums are [-1,0,-1,0]. If we overwrite index 0's sum of -1 with index 2's, we miss the longest valid subarray.

Correct Approach: Only store the first occurrence of each prefix sum value.

2. Forgetting to Initialize with {0: -1}

Starting with an empty hash table misses valid subarrays that begin at index 0:

Incorrect:

sum_to_index = {}  # Missing initial entry

Why this fails: Consider nums = [0,1]. The prefix sum at index 1 is 0, indicating a balanced subarray from the start. Without {0: -1}, we can't detect this because there's no previous occurrence of sum 0 to reference.

Correct Approach: Always initialize with {0: -1} to handle subarrays starting from index 0.

3. Using Wrong Index Calculation

Incorrectly calculating the subarray length by including both endpoints:

Incorrect:

current_length = index - sum_to_index[cumulative_sum] + 1  # WRONG

Why this fails: The subarray spans from sum_to_index[cumulative_sum] + 1 to index (inclusive). The length is simply index - sum_to_index[cumulative_sum], not +1.

Example: If the same sum appears at indices 2 and 5, the balanced subarray is from index 3 to 5, with length = 5 - 2 = 3, not 4.

4. Not Handling Edge Cases

Failing to consider arrays with all 0s or all 1s:

Test these cases:

• nums = [0,0,0] → Should return 0 (no equal distribution possible)
• nums = [1,1] → Should return 0
• nums = [0] → Should return 0 (single element can't have equal 0s and 1s)

The algorithm handles these naturally if implemented correctly, but it's important to verify the logic works for these scenarios.