Problem Statement

Given a binary non empty array, return the maximum length of a contiguous sub-array with an equal number of 0 and 1

Examples

1. For [0,1] longest contiguous array will be [0, 1] itself with length of 2
2. For [1, 1], empty array is the only such contiguous array with count of 0 and 1 equal to 0. So length here will be 0
3. For [0, 0, 0, 1, 1, 0, 1] longest contiguous array will be [0, 0, 1, 1, 0, 1] itself with length of 6

Solution

Points of Interests

Before we look for a solution lets look at what properties of such arrays we can derive from problem statement.

1. Such sub-array will have equal number of 0s and 1s
2. This mean that the sum of such sub-array will be half of their length
3. Or if we consider 0 as -1, then sum of such sub-array will be 0

This whole sum related findings gives us a clue that pre-sum mechanism can be used to solve this problem. With pre-sum in mind, where we consider 0 as -1, we can deduce following

1. For a pre-sum metrics of length N + 1, where pre-sum[i] is sum of all elements from 0th to i-1 th element of the array. Here pre-sum[0] = 0 is just for our convenience.
2. We look closely to a pre-sum metrics that the sum of elements of a sub-array if pre-sum[start - 2] (sum of element before the sub-array) and pre-sum[end - 1] (sum of elements up to end of the sub-array) is equal.
3. So by finding farthest pre-sum values that are equal we can get length of sub-arrays that fits our criteria

Algo

1. Initialize result = 0, a variable where we will store result length of the longest contiguous sub-array
2. Initialize pre_sum array
3. Keep a hash map called first_sum_i, whose key represents the pre-sum value and value represents the first index it appeared at
4. Enumerate over pre_sum
5. ---- If the current sum value from pre_sum value is not in first_sum_i keys, add it with the current index as value
6. ---- If its there, update result in case (first_sum_i[sum] - current index) is greater than result
7. On completing loop we will have desired answer in result variable

Code

def findMaxLength(nums: List[int]) -> int:
    # Initialize pre-sum
    pre_sum = [0]
    for n in nums:
        av = n if n == 1 else -1
        pre_sum.append(pre_sum[-1] + av)
    
    result = 0
    # Stores earliest index of a pre-sum value in pre_sum array
    first_sum_i = dict()
    # Main loop to get largest length
    for i, p in enumerate(pre_sum):
        if p in first_sum_i:
            result = max(result, i - first_sum_i[p])
        else:
            first_sum_i[p] = i
    
    return result