Minimum sum of squared difference

Array Binary Search Greedy Sorting Heap (Priority Queue)

26.8% acceptance

Feb 25, 2026

669

Given nums1, nums2, k1, k2. You may modify elements of nums1 by +1/-1 at most k1 times and elements of nums2 by +1/-1 at most k2 times. Return the minimum sum of squared difference after modifications.

Solution

Rust

Time O(n log n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn min_sum_square_diff(nums1: Vec<i32>, nums2: Vec<i32>, k1: i32, k2: i32) -> i64 {
    let n = nums1.len();
    let mut diffs: Vec<i64> = (0..n)
      .map(|i| (nums1[i] as i64 - nums2[i] as i64).abs())
      .collect();
    let k = k1 as i64 + k2 as i64;
    diffs.sort_unstable_by(|a, b| b.cmp(a)); // descending

    let mut prefix = vec![0i64; n + 1];
    for i in 0..n {
      prefix[i + 1] = prefix[i] + diffs[i];
    }

    // Binary search for minimum target t such that cost(t) <= k
    let mut lo = 0i64;
    let mut hi = diffs[0];
    while lo < hi {
      let mid = lo + (hi - lo) / 2;
      let cnt = diffs.partition_point(|&x| x > mid);
      let cost = prefix[cnt] - mid * cnt as i64;
      if cost <= k {
        hi = mid;
      } else {
        lo = mid + 1;
      }
    }
    let target = lo;
    let cnt = diffs.partition_point(|&x| x > target);
    let remaining = (k - (prefix[cnt] - target * cnt as i64)).min(n as i64) as usize;

    let mut result: i64 = 0;
    for i in 0..n {
      let v = if i < cnt { target } else { diffs[i] };
      let v = if i < remaining { (v - 1).max(0) } else { v };
      result += v * v;
    }
    result
  }
}