LeetCode 3721 - Longest Balanced Subarray II

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Divide and Conquer, Segment Tree, Prefix Sum

Solution

Problem Understanding

The problem requires us to find the length of the longest subarray in an integer array nums where the count of distinct even numbers equals the count of distinct odd numbers. A subarray is a contiguous section of the array, and “distinct” means that duplicate values are counted only once toward the even or odd tally. The input array can have up to 10^5 elements, with each element also up to 10^5.

The output is a single integer, the maximum length of such a balanced subarray. The constraints indicate that a naive approach that examines all possible subarrays (O(n^2) or O(n^3)) would be too slow, so we need an approach that is close to linear or at worst O(n log n).

Edge cases include arrays that are entirely even, entirely odd, have duplicates, or contain just a single element. The problem guarantees the array is non-empty, so we don’t need to handle zero-length input.

In short, the task is to find the longest contiguous section of the array where the number of distinct evens equals the number of distinct odds.

Approaches

The brute-force approach is straightforward: iterate over all possible subarrays using two nested loops. For each subarray, count the number of distinct even and odd numbers using a set. If the counts are equal, update the maximum length found. While correct, this approach has time complexity O(n^3) if implemented naively (or O(n^2) with optimized set operations) and is impractical for n up to 10^5.

The optimal approach relies on the prefix difference technique, inspired by prefix sums. Instead of tracking actual counts of all subarrays, we can encode the difference between the number of distinct evens and distinct odds as we iterate. Specifically, we maintain a running difference diff = count_even_distinct - count_odd_distinct. For each unique configuration, we store the first index where that difference occurs. When the same difference occurs again, the subarray between these indices must be balanced because the difference did not change. To efficiently update distinct counts, we use two hash maps to track last occurrences of each even and odd number. This allows the algorithm to run in O(n) time.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n^3)	O(n)	Iterates all subarrays and counts distinct evens/odds using sets
Optimal	O(n)	O(n)	Uses prefix difference and hash maps to track first occurrences efficiently

Algorithm Walkthrough

Initialize two empty dictionaries, even_count and odd_count, to track the last occurrence of each distinct even and odd number in nums.
Initialize diff_to_index with {0: -1}. This maps the difference diff = #distinct_evens - #distinct_odds to the first index where it occurred.
Initialize diff to 0, which tracks the current difference between distinct even and odd counts.
Initialize max_len to 0, which stores the length of the longest balanced subarray found.
Iterate over nums with index i. For each number num:
If num is even, add it to even_count if it is not already present.
If num is odd, add it to odd_count if it is not already present.
Update diff = len(even_count) - len(odd_count).
If diff is already in diff_to_index, compute the potential balanced subarray length i - diff_to_index[diff] and update max_len if larger.
Otherwise, store diff_to_index[diff] = i.
Return max_len.

Why it works: The key invariant is that if diff is the same at two indices, then the subarray between these indices has the same number of distinct evens and odds added or removed, so it is balanced. Storing the first occurrence ensures the longest subarray for each difference is captured.

Python Solution

from typing import List

class Solution:
    def longestBalanced(self, nums: List[int]) -> int:
        even_count = {}
        odd_count = {}
        diff_to_index = {0: -1}
        max_len = 0

        for i, num in enumerate(nums):
            if num % 2 == 0:
                if num not in even_count:
                    even_count[num] = i
            else:
                if num not in odd_count:
                    odd_count[num] = i

            diff = len(even_count) - len(odd_count)

            if diff in diff_to_index:
                max_len = max(max_len, i - diff_to_index[diff])
            else:
                diff_to_index[diff] = i

        return max_len

This implementation maintains two dictionaries to track distinct numbers for evens and odds. The diff variable encodes the balance, and diff_to_index stores the first occurrence of each diff. By comparing the current index to the first index with the same diff, we efficiently find balanced subarrays.

Go Solution

func longestBalanced(nums []int) int {
    evenCount := make(map[int]bool)
    oddCount := make(map[int]bool)
    diffToIndex := map[int]int{0: -1}
    maxLen := 0

    for i, num := range nums {
        if num%2 == 0 {
            evenCount[num] = true
        } else {
            oddCount[num] = true
        }

        diff := len(evenCount) - len(oddCount)

        if firstIdx, ok := diffToIndex[diff]; ok {
            if i - firstIdx > maxLen {
                maxLen = i - firstIdx
            }
        } else {
            diffToIndex[diff] = i
        }
    }

    return maxLen
}

The Go implementation mirrors the Python logic. Maps track distinct numbers and first occurrences of differences. len(map) is used to count distinct numbers, and the if ok pattern checks for prior occurrence of diff.

Worked Examples

Example 1: nums = [2,5,4,3]

i	num	even_count	odd_count	diff	diff_to_index	max_len
0	2	{2}	{}	1	{0:-1,1:0}	0
1	5	{2}	{5}	0	{0:-1,1:0}	2 (1-(-1))
2	4	{2,4}	{5}	1	{0:-1,1:0}	2
3	3	{2,4}	{5,3}	0	{0:-1,1:0}	4 (3-(-1))

Longest balanced subarray length is 4.

Example 2: nums = [3,2,2,5,4]

i	num	even_count	odd_count	diff	diff_to_index	max_len
0	3	{}	{3}	-1	{-1:0,0:-1}	0
1	2	{2}	{3}	0	{-1:0,0:-1}	2 (1-(-1))
2	2	{2}	{3}	0	{-1:0,0:-1}	3 (2-(-1))
3	5	{2}	{3,5}	-1	{-1:0,0:-1}	3
4	4	{2,4}	{3,5}	0	{-1:0,0:-1}	5 (4-(-1))

Longest balanced subarray length is 5.

Example 3: nums = [1,2,3,2]

i	num	even_count	odd_count	diff	diff_to_index	max_len
0	1	{}	{1}	-1	{-1:0,0:-1}	0
1	2	{2}	{1}	0	{-1:0,0:-1}	2 (1-(-1))
2	3	{2}	{1,3}	-1	{-1:0,0:-1}	2 (2-0)
3	2	{2}	{1,3}	-1	{-1:0,0:-1}	3 (3-0)

Longest balanced subarray length is 3.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each number is processed once, dictionaries provide O(1) operations.
Space	O(n)	Stores distinct numbers and