MediumLeetCode 525Prefix sums + hash map — first occurrence

Contiguous Array (LeetCode 525)

Problem 6 of 10 in Prefix Sums — Practice Problems

The problem

Given a binary array nums (only 0s and 1s), return the length of the longest contiguous subarray that contains the same number of 0s and 1s.

Input:  nums = [0, 1, 0, 1, 1]
Output: 4

The subarray [0, 1, 0, 1] (indices 0 to 3) has two 0s and two 1s. No balanced subarray of length 5 exists.

How to think about it

Start with brute force. Try every start and every end, counting 0s and 1s as you extend. That is O(n²) subarrays. With n up to 10⁵, it is too slow.

First trick: turn “equal counts” into “sum is zero”. Replace each 0 with -1. Now a subarray is balanced exactly when its sum is 0. That turns a counting question into a prefix-sum question.

The key insight: a zero-sum subarray is two equal prefix sums. If the running sum is s at index a and again s at a later index b, then everything in between added up to 0, so nums[a+1..b] is balanced, with length b - a. To make that span as long as possible, compare the current index with the first index where s appeared. So store each running sum the first time you see it and never overwrite it.

Seed the map with {0: -1}. That stands for the empty prefix before index 0, so a balanced subarray starting at index 0 is measured correctly.

Which pattern is this? Prefix sums plus a hash map. The tell is “longest subarray whose sum (or balance) is something”: store the earliest index of each prefix value.

Approach

  1. Set idx = {0: -1}, s = 0, best = 0.
  2. For each index i and value x, add +1 for a 1 and -1 for a 0 to s.
  3. If s is already in idx, the stretch idx[s]+1 .. i is balanced: update best with i - idx[s].
  4. Otherwise record idx[s] = i (the first time this sum appears).
  5. Return best.

Solution

def find_max_length(nums):
    idx = {0: -1}; s = 0; best = 0           # sum 0 "seen" before the array starts
    for i, x in enumerate(nums):
        s += 1 if x == 1 else -1             # 0 counts as -1
        if s in idx: best = max(best, i - idx[s])   # same sum again: balanced stretch
        else: idx[s] = i                     # keep only the first index
    return best
# find_max_length([0,1,0,1,1]) -> 4

Dry run

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

i x s Seen before? best
0 0 -1 no, store idx[-1] = 0 0
1 1 0 yes, at -1 → length 2 2
2 0 -1 yes, at 0 → length 2 2
3 1 0 yes, at -1 → length 4 4
4 1 1 no, store idx[1] = 4 4

The answer is 4.

Complexity

  • Time: O(n). One pass with O(1) average hash-map work per element.
  • Space: O(n). The running sum ranges from -n to n, so the map can hold up to n + 1 keys.

Edge cases and pitfalls

  • No balanced subarray ([1, 1, 1]) returns 0.
  • Forgetting {0: -1} misses subarrays that start at index 0; [0, 1] would wrongly return 0.
  • Overwriting idx[s] with later indices shrinks the span and gives too short an answer. Store only the first occurrence.
  • Counting 0 as 0 instead of -1 breaks everything: the running sum then just counts ones.
  • A sliding window does not work here: the balance can go up and down, so there is no rule for when to shrink.