Minimum operations to equalize subarrays

impl Solution {
  pub fn min_operations(nums: Vec<i32>, k: i32, queries: Vec<Vec<i32>>) -> Vec<i64> {
    let n = nums.len();
    let k64 = k as i64;

    // Prefix sums of nums values for O(1) range sum.
    let mut prefix_sum = vec![0i64; n + 1];
    for i in 0..n {
      prefix_sum[i + 1] = prefix_sum[i] + nums[i] as i64;
    }

    // break_sum[i] = # of boundary positions j in [1..i] where nums[j]%k != nums[j-1]%k.
    // A query [l,r] is consistent iff break_sum[r] == break_sum[l].
    let mut break_sum = vec![0i32; n];
    for i in 1..n {
      break_sum[i] = break_sum[i - 1] + if nums[i] % k != nums[i - 1] % k { 1 } else { 0 };
    }

    // --- Merge Sort Tree (segment tree with sorted arrays + prefix sums) ---
    // Supports range (count_le, sum_le) queries in O(log n) each.
    // Binary-searching for the median over values takes O(log(MAX_VAL) * log n).
    let size = n.next_power_of_two();
    // tree[node] = sorted slice of nums values for that node's range.
    // psum[node][j] = prefix sum of tree[node][0..j].
    let mut tree: Vec<Vec<i32>> = vec![vec![]; 2 * size];
    let mut psum: Vec<Vec<i64>> = vec![vec![]; 2 * size];

    for i in 0..n {
      tree[size + i] = vec![nums[i]];
      psum[size + i] = vec![0, nums[i] as i64];
    }
    for node in (1..size).rev() {
      let (l, r) = (2 * node, 2 * node + 1);
      // Merge two sorted arrays.
      let mut merged = Vec::with_capacity(tree[l].len() + tree[r].len());
      let (mut li, mut ri) = (0, 0);
      while li < tree[l].len() && ri < tree[r].len() {
        if tree[l][li] <= tree[r][ri] {
          merged.push(tree[l][li]);
          li += 1;
        } else {
          merged.push(tree[r][ri]);
          ri += 1;
        }
      }
      merged.extend_from_slice(&tree[l][li..]);
      merged.extend_from_slice(&tree[r][ri..]);
      let mut ps = vec![0i64; merged.len() + 1];
      for (j, &v) in merged.iter().enumerate() {
        ps[j + 1] = ps[j] + v as i64;
      }
      tree[node] = merged;
      psum[node] = ps;
    }

    // Returns (count of elements <= val, their sum) in nums[ql..=qr].
    let query_cs = |ql: usize, qr: usize, val: i32| -> (i64, i64) {
      let mut lo = ql + size;
      let mut hi = qr + size + 1;
      let mut cnt = 0i64;
      let mut sum = 0i64;
      while lo < hi {
        if lo & 1 == 1 {
          let pos = tree[lo].partition_point(|&x| x <= val);
          cnt += pos as i64;
          sum += psum[lo][pos];
          lo += 1;
        }
        if hi & 1 == 1 {
          hi -= 1;
          let pos = tree[hi].partition_point(|&x| x <= val);
          cnt += pos as i64;
          sum += psum[hi][pos];
        }
        lo >>= 1;
        hi >>= 1;
      }
      (cnt, sum)
    };

    queries
      .iter()
      .map(|q| {
        let l = q[0] as usize;
        let r = q[1] as usize;

        // O(1) consistency check.
        if l < r && break_sum[r] != break_sum[l] {
          return -1i64;
        }

        let len = (r - l + 1) as i64;
        let total_sum = prefix_sum[r + 1] - prefix_sum[l];
        // We want the lower median: smallest value v s.t. count_le(v) >= ceil(len/2).
        let target = (len + 1) / 2;

        // Binary search on value space: O(log(1e9) * log n) ≈ 30 * log n.
        let median = {
          let mut lo = 1i32;
          let mut hi = 1_000_000_000i32;
          while lo < hi {
            let mid = lo + (hi - lo) / 2;
            let (cnt, _) = query_cs(l, r, mid);
            if cnt >= target {
              hi = mid;
            } else {
              lo = mid + 1;
            }
          }
          lo
        };

        // sum |x - median| for x in nums[l..=r].
        // = median * cnt_le - sum_le  +  (total_sum - sum_le) - median * (len - cnt_le)
        let (cnt_le, sum_le) = query_cs(l, r, median);
        let cost = median as i64 * cnt_le - sum_le
          + (total_sum - sum_le) - median as i64 * (len - cnt_le);

        // Divide by k (exact because all residues are equal).
        cost / k64
      })
      .collect()
  }
}