LeetCode 3719 - Longest Balanced Subarray I

The problem asks us to find the longest contiguous subarray of an integer array nums such that the number of distinct even numbers in the subarray is equal to the number of distinct odd numbers.

leetcodemediumarrayhash-tabledivide-and-conquersegment-treeprefix-sum

LeetCode Problem 3719

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Divide and Conquer, Segment Tree, Prefix Sum

Solution

Problem Understanding

The problem asks us to find the longest contiguous subarray of an integer array nums such that the number of distinct even numbers in the subarray is equal to the number of distinct odd numbers. In simpler terms, a subarray is "balanced" if it has an equal count of unique even and odd elements, regardless of duplicates.

The input is a list of integers, where 1 <= nums.length <= 1500 and 1 <= nums[i] <= 10^5. The constraints indicate that the array is relatively small, so algorithms with time complexity up to roughly O(n^2) may still be feasible. However, the challenge is to handle the "distinct" property efficiently without counting duplicates multiple times.

The expected output is a single integer representing the length of the longest balanced subarray. Edge cases include arrays with all even or all odd numbers, arrays with only one element, or arrays where multiple balanced subarrays exist but only the longest one should be returned.

Approaches

The brute-force approach iterates through every possible subarray. For each subarray, we track the distinct even and odd numbers using a set, then check if the counts match. If they do, we update the maximum length. This method is guaranteed to give the correct answer because it explicitly checks every subarray, but it is inefficient. Specifically, generating all subarrays takes O(n^2) time, and checking distinct elements within each subarray could take up to O(n) time, giving a total complexity of O(n^3).

The optimal approach uses a prefix-based method. We maintain two hash maps to track the first occurrence of each "even-odd distinct count difference". As we iterate through the array, we update sets of distinct even and odd numbers. For each position, we compute the difference distinct_even_count - distinct_odd_count. If the same difference was seen before, the subarray between the previous occurrence and the current index is balanced. This reduces the time complexity to O(n^2) in practice for our input constraints and avoids re-counting distinct elements repeatedly. Using prefix sets efficiently keeps track of distinct numbers in subarrays.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^3) O(n) Iterate over all subarrays and check distinct counts
Optimal O(n^2) O(n^2) Maintain prefix sets and use difference map for fast lookups

Algorithm Walkthrough

  1. Initialize two empty lists prefix_even and prefix_odd to store sets of distinct even and odd numbers seen up to each index.
  2. Iterate through the array. For each number, if it is even, add it to the last set in prefix_even; if odd, add it to the last set in prefix_odd.
  3. Maintain a dictionary diff_index that maps the difference (len(distinct_even) - len(distinct_odd)) to the earliest index where it occurs.
  4. For each index i, compute curr_diff = len(prefix_even[i]) - len(prefix_odd[i]).
  5. If curr_diff has been seen before at index j, the subarray (j+1, i) is balanced. Update the maximum length if this subarray is longer than the previous maximum.
  6. Return the maximum length found.

Why it works: The key invariant is that whenever the difference between the number of distinct even and odd numbers repeats, the subarray between these two occurrences is balanced. This relies on the idea that the net change in distinct counts between two indices is zero.

Python Solution

from typing import List

class Solution:
    def longestBalanced(self, nums: List[int]) -> int:
        n = len(nums)
        max_len = 0
        prefix_even = [set() for _ in range(n + 1)]
        prefix_odd = [set() for _ in range(n + 1)]
        
        for i in range(n):
            prefix_even[i + 1] = prefix_even[i].copy()
            prefix_odd[i + 1] = prefix_odd[i].copy()
            if nums[i] % 2 == 0:
                prefix_even[i + 1].add(nums[i])
            else:
                prefix_odd[i + 1].add(nums[i])
        
        diff_index = {0: 0}
        for i in range(1, n + 1):
            curr_diff = len(prefix_even[i]) - len(prefix_odd[i])
            if curr_diff in diff_index:
                max_len = max(max_len, i - diff_index[curr_diff])
            else:
                diff_index[curr_diff] = i
        
        return max_len

The implementation first builds prefix sets for distinct even and odd numbers. We then iterate through each index, compute the difference, and use a hash map to efficiently find the longest balanced subarray. Copying sets ensures that distinct counts are correctly maintained up to each index.

Go Solution

func longestBalanced(nums []int) int {
    n := len(nums)
    maxLen := 0
    
    prefixEven := make([]map[int]struct{}, n+1)
    prefixOdd := make([]map[int]struct{}, n+1)
    
    for i := 0; i <= n; i++ {
        prefixEven[i] = make(map[int]struct{})
        prefixOdd[i] = make(map[int]struct{})
    }
    
    for i := 0; i < n; i++ {
        for k, v := range prefixEven[i] {
            prefixEven[i+1][k] = v
        }
        for k, v := range prefixOdd[i] {
            prefixOdd[i+1][k] = v
        }
        if nums[i]%2 == 0 {
            prefixEven[i+1][nums[i]] = struct{}{}
        } else {
            prefixOdd[i+1][nums[i]] = struct{}{}
        }
    }
    
    diffIndex := make(map[int]int)
    diffIndex[0] = 0
    
    for i := 1; i <= n; i++ {
        currDiff := len(prefixEven[i]) - len(prefixOdd[i])
        if prev, ok := diffIndex[currDiff]; ok {
            if i-prev > maxLen {
                maxLen = i - prev
            }
        } else {
            diffIndex[currDiff] = i
        }
    }
    
    return maxLen
}

The Go implementation follows the same logic. Maps replace Python sets, and copying is done manually. The use of struct{} values in the maps is a Go idiomatic way to represent a set.

Worked Examples

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

i prefix_even prefix_odd curr_diff diff_index max_len
0 {} {} 0 {0:0} 0
1 {2} {} 1 {0:0,1:1} 0
2 {2} {5} 0 {0:0,1:1} 2
3 {2,4} {5} 1 {0:0,1:1} 2
4 {2,4} {5,3} 0 {0:0,1:1} 4

Longest balanced subarray length is 4.

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

Following the same process, the longest balanced subarray length is 5.

Complexity Analysis

Measure Complexity Explanation
Time O(n^2) Iterating through the array and copying sets takes O(n^2) in the worst case
Space O(n^2) Storing prefix sets for each index requires O(n^2) space

The algorithm is feasible for n <= 1500, as n^2 operations are around 2.25 million.

Test Cases

# test cases
assert Solution().longestBalanced([2,5,4,3]) == 4 # example 1
assert Solution().longestBalanced([3,2,2,5,4]) == 5 # example 2
assert Solution().longestBalanced([1,2,3,2]) == 3 # example 3
assert Solution().longestBalanced([1,3,5,7]) == 0 # all odd
assert Solution().longestBalanced([2,4,6,8]) == 0 # all even
assert Solution().longestBalanced([1]) == 0 # single element
assert Solution().longestBalanced([2,2,1,1]) == 4 # repeated numbers
assert Solution().longestBalanced([1,2,3,4,5,6,7,8,9,10]) == 10 # alternating even/odd
Test Why
[2,5,4,3] basic small input, balance exists
[3,2,2,5,4] repeated numbers, balance in larger array
[1,2,3,2] subarray in middle is balanced
[1,3