525. Contiguous Array

Problem

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

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 a longest contiguous subarray with equal number of 0 and 1.

Constraints:

• 1 <= nums.length <= 105
• nums[i] is either 0 or 1.

Problem Summary

Given a binary array nums, find the longest contiguous subarray containing the same number of 0s and 1s, and return the length of that subarray.

Solution Approach

• Having the same number of 0s and 1s can be transformed into the difference between their counts being 0. If we treat 0 as -1, then the original problem becomes finding the longest contiguous subarray whose element sum is 0. This again becomes a problem of finding sums within intervals, which naturally can be handled with prefix sums. Suppose the contiguous subarray is the interval [i,j]. The sum of elements in this interval being 0 means prefixSum[j] - prefixSum[i] = 0, that is, prefixSum[i] = prefixSum[j]. Continuously accumulate the prefix sum and store each prefix sum in a map. Once a certain key already exists, it means the prefix sum at some previous index and the current index form an interval whose element sum is 0. This interval is what we are looking for. Scan the entire array, dynamically update the maximum interval length during the scan, and after the scan is complete, the maximum interval length can be obtained, which is the longest contiguous subarray.

Code

package leetcode

func findMaxLength(nums []int) int {
	dict := map[int]int{}
	dict[0] = -1
	count, res := 0, 0
	for i := 0; i < len(nums); i++ {
		if nums[i] == 0 {
			count--
		} else {
			count++
		}
		if idx, ok := dict[count]; ok {
			res = max(res, i-idx)
		} else {
			dict[count] = i
		}
	}
	return res
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}