LeetCode 525 - Contiguous Array

LeetCode Problem 525

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Prefix Sum

Solution

Problem Understanding

This problem asks us to find the maximum length of a contiguous subarray in a binary array where the number of 0s and 1s are equal.

A contiguous subarray means the elements must appear consecutively in the original array. We are not allowed to rearrange elements or skip indices. The goal is to examine all possible contiguous segments and determine which one contains an equal count of 0 and 1, then return the length of the longest such segment.

The input is a binary array nums, meaning every element is either 0 or 1. The output is a single integer representing the length of the longest valid contiguous subarray.

For example:

• In [0,1], the entire array has one 0 and one 1, so the answer is 2.
• In [0,1,0], the full array has two 0s and one 1, which is not balanced. However, [0,1] and [1,0] are balanced, so the maximum length is 2.

The constraints are important:

• 1 <= nums.length <= 10^5
• Each value is either 0 or 1

The input size can be as large as 100,000, which immediately tells us that an O(n²) or O(n³) solution will likely be too slow. We need an algorithm that runs in roughly linear time, ideally O(n).

There are several important edge cases to consider upfront. An array containing only 0s or only 1s can never form a valid balanced subarray, so the result should be 0. Very small arrays, such as length 1, also cannot contain equal numbers of 0 and 1. Another subtle case occurs when the longest valid subarray starts at index 0, which requires careful handling in the implementation to avoid off by one mistakes.

Approaches

Brute Force Approach

A straightforward way to solve this problem is to examine every possible contiguous subarray and count how many 0s and 1s it contains.

We could use two nested loops. The outer loop selects the starting index, and the inner loop expands the ending index. For each subarray, we count the number of 0s and 1s and check whether they are equal. If they are, we update the maximum length.

This approach is correct because it explicitly evaluates every contiguous subarray, guaranteeing that the longest valid one will eventually be found.

However, it is too slow for the given constraints. There are O(n²) subarrays, and counting elements inside each subarray can take O(n) time, resulting in O(n³) complexity. Even with prefix sums to reduce counting to O(1), the algorithm would still require O(n²) time, which is too expensive for n = 10^5.

Optimal Approach, Prefix Sum with Hash Map

The key observation is that we do not actually need to count 0s and 1s separately.

Instead, we can transform the problem:

• Treat 0 as -1
• Treat 1 as +1

Now, whenever a subarray has an equal number of 0s and 1s, its cumulative sum becomes 0.

Why does this work?

Suppose a subarray contains:

• k ones → contributes +k
• k zeros → contributes -k

The total becomes:

(+k) + (-k) = 0

This transforms the problem into finding the longest subarray with sum zero.

We can solve this efficiently using a prefix sum and a hash map.

If the same cumulative sum appears at two different indices, then the elements between those indices must sum to zero.

For example:

Prefix sum at i = 3 → 2
Prefix sum at j = 8 → 2

This means the subarray (i+1 ... j) has sum 0, because:

prefix[j] - prefix[i] = 0

We store the first occurrence of each prefix sum in a hash map. When we see the same sum again, we compute the distance between indices and update the maximum length.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²) or O(n³)	O(1)	Checks every subarray explicitly
Optimal	O(n)	O(n)	Uses prefix sum and hash map to find repeated balances

Algorithm Walkthrough

1. Initialize a hash map to store the first occurrence index of every cumulative balance.

We start with:

{0: -1}

This is crucial because it handles cases where a valid subarray starts from index 0. A balance of 0 before processing any element is conceptually at index -1. 2. Initialize two variables:

• balance = 0
• max_length = 0

The balance tracks the cumulative difference between 1s and 0s. 3. Traverse the array from left to right.

At each index:

• If nums[i] == 1, add 1 to the balance.
• If nums[i] == 0, subtract 1 from the balance.

4. Check whether this balance has been seen before.

If the balance already exists in the hash map, it means the subarray between the previous occurrence and the current index has equal numbers of 0s and 1s.

Compute:

current_length = current_index - first_occurrence_index

Then update the maximum length. 5. If the balance has never been seen before, store its index.

We only store the first occurrence because earlier indices produce longer subarrays. 6. Continue until the entire array has been processed. 7. Return max_length.

Why it works

The algorithm relies on a prefix sum invariant. If the cumulative balance at two indices is the same, then the net contribution of the elements between those indices is zero. Since 0 contributes -1 and 1 contributes +1, a net sum of zero means equal numbers of 0s and 1s. By storing the earliest occurrence of each balance, we maximize the length of valid subarrays.

Python Solution

from typing import List

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
        first_occurrence = {0: -1}
        balance = 0
        max_length = 0

        for index, num in enumerate(nums):
            if num == 1:
                balance += 1
            else:
                balance -= 1

            if balance in first_occurrence:
                current_length = index - first_occurrence[balance]
                max_length = max(max_length, current_length)
            else:
                first_occurrence[balance] = index

        return max_length

The implementation begins by creating a dictionary called first_occurrence, which stores the earliest index where each balance appears. The entry {0: -1} is preloaded so that subarrays beginning at index 0 are handled naturally.

The variable balance tracks the running difference between 1s and 0s. Whenever we encounter a 1, we increase the balance, and whenever we encounter a 0, we decrease it.

For every index, we check whether the current balance has appeared before. If it has, then the segment between the earlier index and the current index contains an equal number of 0s and 1s. We calculate its length and update max_length.

If the balance is new, we store its index. Importantly, we only record the first occurrence because earlier indices produce longer subarrays.

At the end, we return the maximum length found.

Go Solution

func findMaxLength(nums []int) int {
	firstOccurrence := map[int]int{
		0: -1,
	}

	balance := 0
	maxLength := 0

	for index, num := range nums {
		if num == 1 {
			balance++
		} else {
			balance--
		}

		if firstIndex, exists := firstOccurrence[balance]; exists {
			currentLength := index - firstIndex
			if currentLength > maxLength {
				maxLength = currentLength
			}
		} else {
			firstOccurrence[balance] = index
		}
	}

	return maxLength
}

The Go implementation follows the same algorithmic structure as the Python version. A map[int]int is used to store the first occurrence of each balance. Go maps require an explicit existence check using the value, exists := map[key] pattern.

There are no concerns about integer overflow here because the balance can only range between -n and n, which is well within Go's integer limits for n <= 10^5.

Unlike Python, Go does not distinguish between nil and empty slices in a way that matters here, because the problem guarantees at least one element in the input.

Worked Examples

Example 1

Input:

nums = [0,1]

Initial state:

balance = 0
max_length = 0
map = {0: -1}

Index	Value	Balance	Seen Before?	Action	Max Length
0	0	-1	No	Store -1 → 0	0
1	1	0	Yes, at -1	Length = 1 - (-1) = 2	2

Final answer:

2

Example 2

Input:

nums = [0,1,0]

Index	Value	Balance	Seen Before?	Action	Max Length
0	0	-1	No	Store -1 → 0	0
1	1	0	Yes, at -1	Length = 2	2
2	0	-1	Yes, at 0	Length = 2	2

Final answer:

2

Example 3

Input:

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

Index	Value	Balance	Seen Before?	Action	Max Length
0	0	-1	No	Store	0
1	1	0	Yes at -1	Length = 2	2
2	1	1	No	Store	2
3	1	2	No	Store	2
4	1	3	No	Store	2
5	1	4	No	Store	2
6	0	3	Yes at 4	Length = 2	2
7	0	2	Yes at 3	Length = 4	4
8	0	1	Yes at 2	Length = 6	6

Final answer:

6

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	We traverse the array once, and hash map operations are O(1) on average
Space	O(n)	In the worst case, every balance value is unique and stored

The algorithm runs in linear time because each element is processed exactly once. Every hash map lookup and insertion is constant time on average. The extra memory comes from storing prefix balances, which can grow to at most O(n) distinct entries.

Test Cases

solution = Solution()

assert solution.findMaxLength([0, 1]) == 2  # Smallest balanced array
assert solution.findMaxLength([0, 1, 0]) == 2  # Multiple valid subarrays
assert solution.findMaxLength([0, 1, 1, 1, 1, 1, 0, 0, 0]) == 6  # Provided example

assert solution.findMaxLength([0]) == 0  # Single element, no pair
assert solution.findMaxLength([1]) == 0  # Single element, no pair

assert solution.findMaxLength([0, 0, 0]) == 0  # Only zeros
assert solution.findMaxLength([1, 1, 1]) == 0  # Only ones

assert solution.findMaxLength([0, 1, 0, 1]) == 4  # Entire array balanced
assert solution.findMaxLength([1, 0, 1, 0, 1, 0]) == 6  # Alternating values

assert solution.findMaxLength([0, 0, 1, 1]) == 4  # Balanced after imbalance
assert solution.findMaxLength([0, 0, 1, 0, 1, 1, 0]) == 6  # Long subarray in middle

assert solution.findMaxLength([1, 1, 0, 0, 1]) == 4  # Longest subarray not full array
assert solution.findMaxLength([0, 0, 1, 1, 0]) == 4  # Repeated balance values

Test	Why
[0,1]	Smallest valid balanced case
[0,1,0]	Multiple overlapping valid subarrays
[0,1,1,1,1,1,0,0,0]	Provided example with longer segment
[0]	Minimum input size
[1]	Single one, impossible to balance
[0,0,0]	No ones present
[1,1,1]	No zeros present
[0,1,0,1]	Entire array forms answer
[1,0,1,0,1,0]	Alternating values
[0,0,1,1]	Valid subarray appears after imbalance
[0,0,1,0,1,1,0]	Longest segment is internal
[1,1,0,0,1]	Longest subarray excludes suffix
[0,0,1,1,0]	Repeated prefix balances

Edge Cases

Arrays With No Valid Subarray

An array containing only 0s or only 1s cannot possibly contain equal numbers of both values. A naive implementation may accidentally return a non zero value if prefix sums are handled incorrectly. This implementation correctly returns 0 because no repeated balance ever produces a valid length.

Valid Subarray Starting at Index Zero

A common source of bugs is forgetting to handle subarrays that begin at the first index. For example, [0,1] should return 2. Without initializing the hash map with {0: -1}, the algorithm would miss this case. The implementation explicitly includes this setup to ensure correct behavior.

Multiple Occurrences of the Same Balance

Sometimes the same balance appears many times. For example, in [0,1,0,1], balance 0 appears repeatedly. A mistake would be overwriting the earliest index in the hash map, which would reduce the possible subarray length. The implementation avoids this by storing only the first occurrence of each balance, guaranteeing the maximum possible distance is considered.

Very Large Inputs

Since nums.length can reach 100,000, quadratic solutions will time out. This implementation handles large inputs efficiently because it performs only one pass through the array and uses constant time hash map operations, making it scalable to the problem limits.