Separate squares ii

// Segment tree for range-add / covered-length query.
// Leaves represent x-intervals [xs[i], xs[i+1]]; cnt[node] counts how many
// times the node's full range is directly covered by an open rectangle.
// len[node] = total covered length in the node's range.
struct SegTree {
  xs: Vec<i64>,
  cnt: Vec<i32>,
  len: Vec<i64>,
}

impl SegTree {
  fn new(xs: Vec<i64>) -> Self {
    let m = xs.len();
    SegTree { xs, cnt: vec![0; 4 * m], len: vec![0; 4 * m] }
  }

  fn push_up(&mut self, node: usize, lo: usize, hi: usize) {
    if self.cnt[node] > 0 {
      self.len[node] = self.xs[hi] - self.xs[lo];
    } else if hi - lo == 1 {
      self.len[node] = 0;
    } else {
      self.len[node] = self.len[2 * node] + self.len[2 * node + 1];
    }
  }

  // Add `delta` to every leaf interval that lies within [xs[l], xs[r])
  fn update(&mut self, node: usize, lo: usize, hi: usize, l: usize, r: usize, delta: i32) {
    if r <= lo || hi <= l { return; }
    if l <= lo && hi <= r {
      self.cnt[node] += delta;
      self.push_up(node, lo, hi);
      return;
    }
    let mid = (lo + hi) / 2;
    self.update(2 * node, lo, mid, l, r, delta);
    self.update(2 * node + 1, mid, hi, l, r, delta);
    self.push_up(node, lo, hi);
  }
}

impl Solution {
  pub fn separate_squares(squares: Vec<Vec<i32>>) -> f64 {
    // 1. Compress x-coordinates.
    let mut xs: Vec<i64> = squares
      .iter()
      .flat_map(|s| [s[0] as i64, s[0] as i64 + s[2] as i64])
      .collect();
    xs.sort_unstable();
    xs.dedup();
    let m = xs.len();
    if m < 2 { return 0.0; }
    let n_intervals = m - 1;

    // 2. Build events: (y, delta, x_left_idx, x_right_idx)
    //    +1 when a square starts, -1 when it ends.
    let mut events: Vec<(i64, i32, usize, usize)> = Vec::with_capacity(squares.len() * 2);
    for s in &squares {
      let xi = s[0] as i64;
      let yi = s[1] as i64;
      let li = s[2] as i64;
      let il = xs.partition_point(|&v| v < xi);
      let ir = xs.partition_point(|&v| v < xi + li);
      events.push((yi,      1, il, ir));
      events.push((yi + li, -1, il, ir));
    }
    // Sort by y; at the same y process additions (+1) before removals (-1)
    // so that a square ending exactly where another starts is handled correctly.
    events.sort_unstable_by(|a, b| a.0.cmp(&b.0).then(b.1.cmp(&a.1)));

    // 3. One-pass sweep: accumulate (y_start, dy, covered_x_width) strips.
    let mut seg = SegTree::new(xs);
    let mut strips: Vec<(i64, i64, i64)> = Vec::new(); // (y_start, dy, width)

    let mut ei = 0;
    let ne = events.len();
    let mut prev_y = events[0].0;

    while ei < ne {
      let cur_y = events[ei].0;
      if cur_y > prev_y {
        strips.push((prev_y, cur_y - prev_y, seg.len[1]));
      }
      prev_y = cur_y;
      while ei < ne && events[ei].0 == cur_y {
        let (_, delta, il, ir) = events[ei];
        seg.update(1, 0, n_intervals, il, ir, delta);
        ei += 1;
      }
    }

    // 4. Find the y that splits total area in half.
    //    Work in integer arithmetic (×2) to avoid precision loss until the final step.
    let total: i128 = strips.iter().map(|&(_, dy, w)| dy as i128 * w as i128).sum();
    let half_x2 = total; // 2 * half

    let mut cum2: i128 = 0;
    for &(y_start, dy, width) in &strips {
      let area2 = 2 * dy as i128 * width as i128;
      if cum2 + area2 >= half_x2 {
        if width == 0 {
          return y_start as f64;
        }
        // y_start + (half - cum) / width
        // = y_start + (half_x2 - cum2) / (2 * width)
        let remaining = (half_x2 - cum2) as f64;
        return y_start as f64 + remaining / (2.0 * width as f64);
      }
      cum2 += area2;
    }

    // Fallback (unreachable for valid input)
    prev_y as f64
  }
}