Longest balanced subarray ii

// O(n log n) approach using a segment tree with range-add and leftmost-zero query.
//
// Key observation: define D[l] = distinct_evens(l, r) - distinct_odds(l, r) for the
// current right endpoint r. We want the leftmost l in [0, r] with D[l] = 0.
//
// As r advances by one (processing nums[r] = v with parity delta ∈ {+1,−1}):
//   let p = last occurrence of v before r (−1 if none)
//   range-add delta to D[p+1 .. r]   (v becomes newly distinct for any window starting in [p+1,r])
//
// Initial state: all D[l] = 0 (empty window is balanced).
// Segment tree stores (min, max) per node with lazy range-add.
// Leftmost-zero query in O(log n): descend into left child first whenever that child's
// range [min, max] brackets 0; only visit right child if left yields nothing.

impl Solution {
  pub fn longest_balanced(nums: Vec<i32>) -> i32 {
    let n = nums.len();
    if n == 0 {
      return 0;
    }
    // Segment tree arrays (1-indexed, size 4n)
    let sz = 4 * n;
    let mut mn = vec![0i32; sz];  // min value in range
    let mut mx = vec![0i32; sz];  // max value in range
    let mut lazy = vec![0i32; sz];

    let mut prev: std::collections::HashMap<i32, i32> = std::collections::HashMap::new();
    let mut best = 0i32;

    for r in 0..n {
      let v = nums[r];
      let p = *prev.get(&v).unwrap_or(&-1);
      let lo = (p + 1) as usize;
      let delta = if v % 2 == 0 { 1i32 } else { -1i32 };
      Self::update(&mut mn, &mut mx, &mut lazy, 1, 0, n - 1, lo, r, delta);
      prev.insert(v, r as i32);

      // Find leftmost l in [0, r] with D[l] == 0
      if let Some(l) = Self::find_zero(&mut mn, &mut mx, &mut lazy, 1, 0, n - 1, 0, r) {
        best = best.max((r - l + 1) as i32);
      }
    }
    best
  }

  fn push_down(mn: &mut Vec<i32>, mx: &mut Vec<i32>, lazy: &mut Vec<i32>, node: usize) {
    let d = lazy[node];
    if d != 0 {
      for c in [2 * node, 2 * node + 1] {
        mn[c] += d;
        mx[c] += d;
        lazy[c] += d;
      }
      lazy[node] = 0;
    }
  }

  fn update(
    mn: &mut Vec<i32>, mx: &mut Vec<i32>, lazy: &mut Vec<i32>,
    node: usize, lo: usize, hi: usize,
    l: usize, r: usize, delta: i32,
  ) {
    if l > hi || r < lo {
      return;
    }
    if l <= lo && hi <= r {
      mn[node] += delta;
      mx[node] += delta;
      lazy[node] += delta;
      return;
    }
    Self::push_down(mn, mx, lazy, node);
    let mid = (lo + hi) / 2;
    Self::update(mn, mx, lazy, 2 * node, lo, mid, l, r, delta);
    Self::update(mn, mx, lazy, 2 * node + 1, mid + 1, hi, l, r, delta);
    mn[node] = mn[2 * node].min(mn[2 * node + 1]);
    mx[node] = mx[2 * node].max(mx[2 * node + 1]);
  }

  /// Returns the leftmost index in [ql, qr] where D == 0, or None.
  fn find_zero(
    mn: &mut Vec<i32>, mx: &mut Vec<i32>, lazy: &mut Vec<i32>,
    node: usize, lo: usize, hi: usize,
    ql: usize, qr: usize,
  ) -> Option<usize> {
    if lo > qr || hi < ql { return None; }
    // No zero possible in this node's range
    if mn[node] > 0 || mx[node] < 0 { return None; }
    if lo == hi { return Some(lo); }
    Self::push_down(mn, mx, lazy, node);
    let mid = (lo + hi) / 2;
    let left = Self::find_zero(mn, mx, lazy, 2 * node, lo, mid, ql, qr);
    if left.is_some() {
      return left;
    }
    Self::find_zero(mn, mx, lazy, 2 * node + 1, mid + 1, hi, ql, qr)
  }
}