525. Contiguous Array

LeetCode 525. Contiguous Array

Description

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

Example 1:

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

Example 2:

Input: [0,1,0]
Output: 2
Explanation: [0, 1] (or [1, 0]) is a longest contiguous subarray with equal number of 0 and 1.

Note: The length of the given binary array will not exceed 50,000.

Tags

Hash Table

Solution

This problem is similar to LeetCode 523. Continuous Subarray Sum, and we are going to apply the prefix-sum stategy as well. The requirement is equivalent to find the longest subarray that contains an equal number of 1s and 0s. If we replace all 0s to -1, the sum of the target subarray is 0. Thus, we can use a prefix-sum array to compute sum values of all subarrays, and pick the longest one whose sum is 0.

In fact, we do not need such prefix-sum array. Instead, we can use a counter variable to count 1s and 0s. When iterating over the input array, if we meet 1, then we increase the counter by 1; otherwise, we decrease it by 1. Also, we use a hash table that maps a first occurrence counter value to the index, initialized with {0: -1} representing an empty array. If a counter value appears twice, then we can find a subarray who starts from previous occurrence index and end with the current one. Record the length of this subarray if it is longer.

Complexity

• Time complexity: $O(n)$

• Space complexity: $O(n)$

Code

func findMaxLength(nums []int) int {
	ans, count, cnt2idx := 0, 0, map[int]int{0: -1}
	for i, num := range nums {
		if num == 1 {
			count++
		} else {
			count--
		}
		if preIdx, ok := cnt2idx[count]; ok {
			if ans < i-preIdx {
				ans = i - preIdx
			}
		} else {
			cnt2idx[count] = i
		}
	}
	return ans
}

Reference

1. 连续数组