Minimum operations to make elements within k subarrays equal

impl Solution {
  pub fn min_operations(nums: Vec<i32>, x: i32, k: i32) -> i64 {
    let x = x as usize;
    let k = k as usize;

    // Compute minimum-cost window for every window of size x using a
    // coordinate-compressed Fenwick tree (BIT).  This runs in O(n log n)
    // compared to the naive O(n * x) approach with BTreeMap.
    let wcost = window_costs(&nums, x);
    let nw = wcost.len(); // = n - x + 1

    // DP: dp[j][i] = min total cost when j-th (1-indexed) non-overlapping
    // window is placed at wcost index i.
    // Transition: dp[j][i] = wcost[i] + min_{p : p + x <= i} dp[j-1][p]
    // A rolling prefix-min reduces this to O(k * n).
    const INF: i64 = i64::MAX / 2;

    let mut prev = wcost.clone(); // layer j = 1

    for _layer in 1..k {
      let mut cur = vec![INF; nw];
      let mut pm = INF; // min of prev[0 .. i-x]
      for i in 0..nw {
        if i >= x {
          let v = prev[i - x];
          if v < pm {
            pm = v;
          }
        }
        if pm < INF {
          let candidate = wcost[i].saturating_add(pm);
          if candidate < cur[i] {
            cur[i] = candidate;
          }
        }
      }
      prev = cur;
    }

    *prev.iter().filter(|&&v| v < INF).min().unwrap_or(&INF)
  }
}

// ---------------------------------------------------------------------------
// Fenwick tree (BIT) supporting point-update, prefix-query, and O(log n)
// k-th order statistic via binary lifting.
// ---------------------------------------------------------------------------
struct BIT {
  n: usize,
  tree: Vec<i64>,
}

impl BIT {
  fn new(n: usize) -> Self {
    BIT { n, tree: vec![0i64; n + 1] }
  }

  // Add `delta` to position `i` (0-indexed).
  fn update(&mut self, i: usize, delta: i64) {
    let mut j = i + 1; // 1-indexed internally
    while j <= self.n {
      self.tree[j] += delta;
      j += j & j.wrapping_neg();
    }
  }

  // Prefix sum [0 .. i] inclusive (0-indexed i).
  fn query(&self, i: usize) -> i64 {
    let mut j = i + 1;
    let mut s = 0i64;
    while j > 0 {
      s += self.tree[j];
      j -= j & j.wrapping_neg();
    }
    s
  }

  // 0-indexed position of the k-th smallest element (k is 1-indexed).
  // Requires all stored counts to be non-negative and total count >= k.
  fn kth(&self, mut k: i64) -> usize {
    let mut pos = 0usize;
    let log = usize::BITS - self.n.leading_zeros();
    let mut bit = 1usize << (log - 1);
    while bit > 0 {
      let nxt = pos + bit;
      if nxt <= self.n && self.tree[nxt] < k {
        k -= self.tree[nxt];
        pos = nxt;
      }
      bit >>= 1;
    }
    pos
  }
}

// ---------------------------------------------------------------------------
// For each sliding window of size x compute sum(|elem - median|).
// Uses coordinate compression + two BITs (count, sum) → O(n log n) total.
// ---------------------------------------------------------------------------
fn window_costs(nums: &[i32], x: usize) -> Vec<i64> {
  let n = nums.len();
  let nw = n - x + 1;

  // Coordinate compression
  let mut sorted_vals: Vec<i64> = nums.iter().map(|&v| v as i64).collect();
  sorted_vals.sort_unstable();
  sorted_vals.dedup();
  let m = sorted_vals.len();

  let compress = |v: i64| -> usize { sorted_vals.partition_point(|&x| x < v) };

  let mut cnt_bit = BIT::new(m);
  let mut sum_bit = BIT::new(m);
  let mut total_sum = 0i64;

  let mut costs = vec![0i64; nw];

  for i in 0..n {
    let v = nums[i] as i64;
    let ci = compress(v);
    cnt_bit.update(ci, 1);
    sum_bit.update(ci, v);
    total_sum += v;

    if i >= x {
      let rv = nums[i - x] as i64;
      let rci = compress(rv);
      cnt_bit.update(rci, -1);
      sum_bit.update(rci, -rv);
      total_sum -= rv;
    }

    if i >= x - 1 {
      // Median = (lo_target + 1)-th smallest element (1-indexed).
      // lo_target elements lie strictly "below" the median slot.
      let lo_target = x / 2;
      let med_ci = cnt_bit.kth(lo_target as i64 + 1);
      let med = sorted_vals[med_ci];

      let (cnt_below, sum_below) = if med_ci > 0 {
        (cnt_bit.query(med_ci - 1), sum_bit.query(med_ci - 1))
      } else {
        (0, 0)
      };

      let sum_above = total_sum - sum_bit.query(med_ci);
      let cnt_above = (x as i64) - cnt_bit.query(med_ci);

      // cost = Σ|elem - med|
      //      = med * cnt_below - sum_below   (lower half)
      //      + sum_above - med * cnt_above   (upper half)
      costs[i + 1 - x] = med * cnt_below - sum_below + sum_above - med * cnt_above;
    }
  }

  costs
}