3721. Longest Balanced Subarray II

  • Array
  • Hash Table
  • Segment Tree
TIME
O(n log n)
SPACE
O(n)
problem on leetcode ↗source on github ↗

Intuition

For each left boundary l, we need to find the rightmost r where distinctOdd[l..r] == distinctEven[l..r]. Define diff[r] = distinctOdd - distinctEven for subarray [l..r]. We need to find rightmost r where diff[r] == 0.

Solution

  1. Transform to balance array: For each position, track balance = (distinct odds) - (distinct evens). A subarray is balanced when balance = 0.

  2. Segment tree with lazy propagation:

    • Track min and max in ranges to quickly find if 0 exists
    • Use lazy propagation for efficient range updates

  3. Sliding window on left boundary:

    • For each l, find rightmost r where balance = 0 using segment tree
    • When moving l to l+1, update the contribution of nums[l]

  4. First occurrence tracking:

    • For each value, track positions in a deque
    • When a value's first occurrence moves, update the segment tree range

Key Insight

When left boundary moves from l to l+1:

  • If nums[l] was the first occurrence of that value in [l, n-1]
  • Its contribution (+1 for odd, -1 for even) needs to be removed from range [l, next_occurrence-1]

Complexity

  • Time: O(n log n), each position processed once, segment tree operations are O(log n)
  • Space: O(n), segment tree and position tracking

Edge Cases

  • All even or all odd → return 0
  • Single element → return 0
  • All same value → return 0

Code

from typing import List
from collections import defaultdict, deque


class SegmentTree:
    def __init__(self, n):
        self.n = n
        self.minv = [0] * (4 * n)
        self.maxv = [0] * (4 * n)
        self.lazy = [0] * (4 * n)

    def _push(self, idx):
        if self.lazy[idx] != 0:
            for child in (idx * 2, idx * 2 + 1):
                self.minv[child] += self.lazy[idx]
                self.maxv[child] += self.lazy[idx]
                self.lazy[child] += self.lazy[idx]
            self.lazy[idx] = 0

    def _pull(self, idx):
        self.minv[idx] = min(self.minv[idx * 2], self.minv[idx * 2 + 1])
        self.maxv[idx] = max(self.maxv[idx * 2], self.maxv[idx * 2 + 1])

    def range_add(self, l, r, val, idx=1, nl=0, nr=None):
        if nr is None:
            nr = self.n - 1
        if r < nl or nr < l:
            return
        if l <= nl and nr <= r:
            self.minv[idx] += val
            self.maxv[idx] += val
            self.lazy[idx] += val
            return
        self._push(idx)
        mid = (nl + nr) // 2
        self.range_add(l, r, val, idx * 2, nl, mid)
        self.range_add(l, r, val, idx * 2 + 1, mid + 1, nr)
        self._pull(idx)

    def find_zero(self, ql, idx=1, nl=0, nr=None):
        if nr is None:
            nr = self.n - 1
        if nr < ql or self.minv[idx] > 0 or self.maxv[idx] < 0:
            return -1
        if nl == nr:
            return nl
        self._push(idx)
        mid = (nl + nr) // 2
        res = self.find_zero(ql, idx * 2 + 1, mid + 1, nr)
        if res != -1:
            return res
        return self.find_zero(ql, idx * 2, nl, mid)


class Solution:
    def longestBalanced(self, nums: List[int]) -> int:
        """
        3721. Longest Balanced Subarray II
        Time: O(n log n)
        Space: O(n)
        """
        n = len(nums)
        pos = defaultdict(deque)

        for i, v in enumerate(nums):
            pos[v].append(i)

        st = SegmentTree(n)

        for v, dq in pos.items():
            sign = 1 if v % 2 else -1
            st.range_add(dq[0], n - 1, sign)

        ans = 0

        for l in range(n):
            r = st.find_zero(l)
            if r != -1:
                ans = max(ans, r - l + 1)

            v = nums[l]
            sign = 1 if v % 2 else -1
            dq = pos[v]
            dq.popleft()
            nxt = dq[0] if dq else n
            st.range_add(l, nxt - 1, -sign)

        return ans