Given a binary array nums, return the maximum length of a contiguous subarray with an equal number of 0 and 1.

Examples

Example 1:

Input: nums = [0,1]
Output: 2
Explanation: [0, 1] is the longest contiguous subarray with an equal number of 0 and 1.

Example 2:

Input: nums = [0,1,0]
Output: 2
Explanation: [0, 1] or [1, 0] is the longest contiguous subarray with an equal number of 0 and 1.

Constraints

  • 1 <= nums.length <= 10^5
  • nums[i] is either 0 or 1.

Thinking Process

  1. Transformation: 0 → -1, 1 → +1 converts the problem to finding subarrays with sum 0
  • Clarify if the array is sorted, has negatives, or allows duplicates.
  • Prefix sums answer range queries; hash maps answer pair/count queries.
  • In-place tricks use swap/write index instead of extra arrays.
Array + hash map 2 7 11 map hash map for O(1) lookups

Common Approaches

Typical techniques for this pattern:

Approach Time Space Notes
Prefix sum (this problem) O(n) O(n) Range queries, subarray sum
Sort + scan O(n log n) O(1) Intervals, meeting rooms
Kadane’s algorithm O(n) O(1) Maximum subarray
Hash map counting O(n) O(n) Frequency, two-sum variants

Solution

Time Complexity: O(n)
Space Complexity: O(n)

Convert 0 to -1 and 1 to +1, then use prefix sum. When we see a prefix sum we’ve encountered before, the subarray between those two positions has equal 0s and 1s (net sum of 0).

class Solution:
    def findMaxLength(self, nums):
        first_index = {0: -1}  # prefix sum 0 at virtual index -1

        max_len = 0
        count = 0

        for i in range(len(nums)):
            count += 1 if nums[i] == 1 else -1

            if count in first_index:
                max_len = max(max_len, i - first_index[count])
            else:
                first_index[count] = i

        return max_len

Solution Explanation

Approach: Prefix sum (this problem)

Key idea: 1. Transformation: 0 → -1, 1 → +1 converts the problem to finding subarrays with sum 0

How the code works:

  1. Transformation: 0 → -1, 1 → +1 converts the problem to finding subarrays with sum 0
    • Clarify if the array is sorted, has negatives, or allows duplicates.
    • Prefix sums answer range queries; hash maps answer pair/count queries.
    • In-place tricks use swap/write index instead of extra arrays.

Walkthrough — input nums = [0,1], expected output 2:

[0, 1] is the longest contiguous subarray with an equal number of 0 and 1.

| Aspect | Complexity | |——–|————| | Time | O(n) - Single pass through array, hash map operations are O(1) | | Space | O(n) - Hash map stores at most n prefix sums |

Algorithm Breakdown

1. Initialize Hash Map

first_index: dict[int, int] = {}
first_index[0] = -1  # Base case
max_len, count = 0, 0
  • map[0] = -1: Represents prefix sum 0 before the array starts
  • This allows us to handle subarrays starting at index 0
  • count: Tracks the current prefix sum

2. Transform and Calculate Prefix Sum

count += 1 if nums[i] == 1 else -1
  • Convert: 0 → -1, 1 → +1
  • Accumulate: Add to running count

3. Check for Repeated Prefix Sum

if count in first_index:
    max_len = max(max_len, i - first_index[count])
else:
    first_index[count] = i
  • If prefix sum seen before: Subarray from map[count] + 1 to i has sum 0
  • Length: i - map[count] (we don’t add 1 because map[count] is the index before the subarray starts)

4. Store First Occurrence

def find_max_length_bruteforce(nums: list[int]) -> int:
    max_len = 0
    for i in range(len(nums)):
        zeros = ones = 0
        for j in range(i, len(nums)):
            if nums[j] == 0:
                zeros += 1
            else:
                ones += 1
            if zeros == ones:
                max_len = max(max_len, j - i + 1)
    return max_len
  • Store only first occurrence: To maximize subarray length
  • Later occurrences: Would give shorter subarrays

Why This Works

Mathematical Proof

For a subarray nums[i+1..j] to have equal 0s and 1s:

  • Number of 1s = Number of 0s
  • Sum of transformed values = 0
  • prefix[j] - prefix[i] = 0
  • prefix[j] = prefix[i]

Therefore, if prefix[j] == prefix[i], the subarray nums[i+1..j] has equal 0s and 1s.

Example Explanation

nums = [0, 1, 0, 0, 1, 1, 0]
       0  1  2  3  4  5  6

Transformed: [-1, +1, -1, -1, +1, +1, -1]

Prefix sums:
i=-1: prefix = 0 (base case)
i=0:  prefix = -1
i=1:  prefix = 0  ← Same as i=-1!
i=2:  prefix = -1 ← Same as i=0!
i=3:  prefix = -2
i=4:  prefix = -1 ← Same as i=0!
i=5:  prefix = 0   ← Same as i=-1!
i=6:  prefix = -1   ← Same as i=0!

Subarray [0,5]: prefix[5] = 0 = prefix[-1]
  Length = 5 - (-1) = 6 ✓
  
Subarray [1,5]: prefix[5] = 0 = prefix[1]
  Length = 5 - 1 = 4 ✓

Complexity

| Aspect | Complexity | |——–|————| | Time | O(n) - Single pass through array, hash map operations are O(1) | | Space | O(n) - Hash map stores at most n prefix sums |

Common Mistakes

  1. All zeros or all ones: [0,0,0] or [1,1,1] → 0 (no equal subarray)
  2. Single element: [0] or [1] → 0
  3. Perfect balance: [0,1] → 2

  4. Multiple valid subarrays: [0,1,0,1,0,1] → 6

  5. Missing base case: Forgetting map[0] = -1 causes issues with subarrays starting at index 0
  6. Wrong length calculation: Using i - map[count] + 1 instead of i - map[count]
  7. Updating map incorrectly: Updating map even when count exists (should only store first occurrence)
  8. Wrong transformation: Using 0 → 0 and 1 → 1 instead of 0 → -1 and 1 → +1

Detailed Example Walkthrough

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

Step 0: Initialize
  map = {0: -1}
  count = 0
  maxLen = 0

Step 1: i=0, nums[0]=0
  count = 0 + (-1) = -1
  map.contains(-1)? No
  map[-1] = 0
  map = {0: -1, -1: 0}
  maxLen = 0

Step 2: i=1, nums[1]=1
  count = -1 + 1 = 0
  map.contains(0)? Yes (at index -1)
  maxLen = max(0, 1 - (-1)) = max(0, 2) = 2
  map = {0: -1, -1: 0}
  maxLen = 2

Step 3: i=2, nums[2]=0
  count = 0 + (-1) = -1
  map.contains(-1)? Yes (at index 0)
  maxLen = max(2, 2 - 0) = max(2, 2) = 2
  map = {0: -1, -1: 0}
  maxLen = 2

Step 4: i=3, nums[3]=0
  count = -1 + (-1) = -2
  map.contains(-2)? No
  map[-2] = 3
  map = {0: -1, -1: 0, -2: 3}
  maxLen = 2

Step 5: i=4, nums[4]=1
  count = -2 + 1 = -1
  map.contains(-1)? Yes (at index 0)
  maxLen = max(2, 4 - 0) = max(2, 4) = 4
  map = {0: -1, -1: 0, -2: 3}
  maxLen = 4

Step 6: i=5, nums[5]=1
  count = -1 + 1 = 0
  map.contains(0)? Yes (at index -1)
  maxLen = max(4, 5 - (-1)) = max(4, 6) = 6
  map = {0: -1, -1: 0, -2: 3}
  maxLen = 6

Step 7: i=6, nums[6]=0
  count = 0 + (-1) = -1
  map.contains(-1)? Yes (at index 0)
  maxLen = max(6, 6 - 0) = max(6, 6) = 6
  map = {0: -1, -1: 0, -2: 3}
  maxLen = 6

Final result: 6

Why Store Only First Occurrence?

To maximize subarray length, we want the earliest starting position for each prefix sum.

Example:

prefix sums: [0, -1, 0, -1, 0, -1]
              ↑   ↑   ↑   ↑   ↑   ↑
            i=-1 0   1   2   3   4

For prefix sum 0:
- First occurrence: i = -1
- Later occurrence: i = 1

Subarray ending at i=3:
- Using i=-1: length = 3 - (-1) = 4 ✓
- Using i=1:  length = 3 - 1 = 2 ✗

Therefore, we should store only the first occurrence.

Pattern Recognition

This problem demonstrates the Prefix Sum + Hash Map pattern:

  • Transform problem into finding subarrays with target sum
  • Use prefix sums to calculate subarray sums efficiently
  • Hash map stores first occurrence of each prefix sum
  • Look for repeated prefix sums to find valid subarrays

Key Insight:

  • Subarray nums[i+1..j] has sum 0 if prefix[j] == prefix[i]
  • This extends to any target sum: subarray has sum k if prefix[j] - prefix[i] == k

Applications:

  • Finding subarrays with specific sum
  • Finding subarrays with equal elements
  • Finding subarrays with specific properties (balanced, etc.)

Optimization Tips

Early Exit (Not Applicable)

Since we need to check all positions, early exit isn’t possible.

Memory Optimization

For very large arrays, consider using array instead of hash map if prefix sum range is bounded.

Use contains() vs find()

def find_max_length_array(nums: list[int]) -> int:
    n = len(nums)
    offset = n
    first_at = [-2] * (2 * n + 1)  # -2 = not seen; prefix 0 stored at index offset
    first_at[offset] = -1
    max_len = count = 0
    for i in range(n):
        count += 1 if nums[i] == 1 else -1
        idx = count + offset
        if first_at[idx] != -2:
            max_len = max(max_len, i - first_at[idx])
        else:
            first_at[idx] = i
    return max_len

Both are O(1) average case, but contains() is more readable.

Code Quality Notes

  1. Readability: Clear variable names and logic
  2. Efficiency: Optimal O(n) time and space
  3. Correctness: Handles all edge cases properly
  4. Modern C++: Uses contains() method (Python20)

This problem is a classic example of transforming a problem into a prefix sum problem. The key insight is converting the “equal 0s and 1s” constraint into “sum equals 0” after transformation.

Key Takeaways

  1. Transformation: 0 → -1, 1 → +1 converts the problem to finding subarrays with sum 0
  2. Prefix Sum: Track cumulative count from start
  3. Hash Map: Store first occurrence of each prefix sum for longest subarray
  4. Base Case: map[0] = -1 handles subarrays starting from index 0
  5. Repeated Prefix: If prefix[i] == prefix[j], then subarray[i+1..j] has sum 0

References

Template Reference