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: 4The 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
- Set
idx = {0: -1},s = 0,best = 0. - For each index
iand valuex, add+1for a 1 and-1for a 0 tos. - If
sis already inidx, the stretchidx[s]+1 .. iis balanced: updatebestwithi - idx[s]. - Otherwise record
idx[s] = i(the first time this sum appears). - 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]) -> 4Dry 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.
Related problems
- Maximum Size Subarray Sum Equals k (LeetCode 325): the same first-index map, for any target sum.
- Subarray Sum Equals K (LeetCode 560): count the subarrays instead of finding the longest.
- Continuous Subarray Sum (LeetCode 523): first-index map on remainders instead of sums.
- 1.Range Sum Query - Immutable (LeetCode 303)
- 2.Find Pivot Index (LeetCode 724)
- 3.Minimum Value to Get Positive Step by Step Sum (LeetCode 1413)
- 4.Subarray Sum Equals K (LeetCode 560)
- 5.Product of Array Except Self (LeetCode 238)
- 6.Contiguous Array (LeetCode 525)
- 7.Subarray Sums Divisible by K (LeetCode 974)
- 8.Continuous Subarray Sum (LeetCode 523)
- 9.Range Sum Query 2D - Immutable (LeetCode 304)
- 10.Maximum Size Subarray Sum Equals k (LeetCode 325)