Minimum operations to make subarray elements equal

Array Hash Table Math Sliding Window Heap (Priority Queue)

46.6% acceptance

Mar 31, 2026

You are given an integer array nums and an integer k. You can perform the following operation any number of times: Increase or decrease any element of nums by 1. Return the minimum number of operations required to ensure that at least one subarray of size k in nums has all elements equal.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn min_operations(nums: Vec<i32>, k: i32) -> i64 {
    let k = k as usize;
    let n = nums.len();
    let offset = 1_000_001i64;
    let sz = 2_000_002usize;

    let mut cnt = vec![0i64; sz + 1];
    let mut sm = vec![0i64; sz + 1];

    fn upd(tree: &mut [i64], mut i: usize, val: i64) {
      while i < tree.len() {
        tree[i] += val;
        i += i & i.wrapping_neg();
      }
    }

    fn qry(tree: &[i64], mut i: usize) -> i64 {
      let mut s = 0;
      while i > 0 {
        s += tree[i];
        i -= i & i.wrapping_neg();
      }
      s
    }

    fn kth(cnt_tree: &[i64], k: i64, sz: usize) -> usize {
      let mut pos = 0;
      let mut rem = k;
      let mut bit = 1;
      while bit <= sz { bit <<= 1; }
      bit >>= 1;
      while bit > 0 {
        let next = pos + bit;
        if next <= sz && cnt_tree[next] < rem {
          rem -= cnt_tree[next];
          pos = next;
        }
        bit >>= 1;
      }
      pos + 1
    }

    let mut ans = i64::MAX;

    for i in 0..n {
      let v = (nums[i] as i64 + offset) as usize;
      upd(&mut cnt, v, 1);
      upd(&mut sm, v, nums[i] as i64);

      if i >= k {
        let old = (nums[i - k] as i64 + offset) as usize;
        upd(&mut cnt, old, -1);
        upd(&mut sm, old, -(nums[i - k] as i64));
      }

      if i + 1 >= k {
        let med_idx = kth(&cnt, ((k + 1) / 2) as i64, sz);
        let median = med_idx as i64 - offset;

        let cnt_le = qry(&cnt, med_idx);
        let sum_le = qry(&sm, med_idx);

        let cnt_gt = k as i64 - cnt_le;
        let sum_gt = qry(&sm, sz) - sum_le;

        let cost = median * cnt_le - sum_le + sum_gt - median * cnt_gt;
        ans = ans.min(cost);
      }
    }

    ans
  }
}