Add problem 2602: Minimum Operations to Make All Array Elements Equal

Author: EFanZh

src/lib.rs

@@ -1925,6 +1925,7 @@ pub mod problem_2596_check_knight_tour_configuration;
 pub mod problem_2598_smallest_missing_non_negative_integer_after_operations;
 pub mod problem_2600_k_items_with_the_maximum_sum;
 pub mod problem_2601_prime_subtraction_operation;
+pub mod problem_2602_minimum_operations_to_make_all_array_elements_equal;
 
 #[cfg(test)]
 mod test_utilities;

src/problem_2602_minimum_operations_to_make_all_array_elements_equal/mod.rs

@@ -0,0 +1,21 @@
+pub mod prefix_sums_and_binary_search;
+
+pub trait Solution {
+    fn min_operations(nums: Vec<i32>, queries: Vec<i32>) -> Vec<i64>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((&[3, 1, 6, 8] as &[_], &[1, 5] as &[_]), &[14_i64, 10] as &[_]),
+            ((&[2, 9, 6, 3], &[10]), &[20]),
+        ];
+
+        for ((nums, queries), expected) in test_cases {
+            assert_eq!(S::min_operations(nums.to_vec(), queries.to_vec()), expected);
+        }
+    }
+}

src/problem_2602_minimum_operations_to_make_all_array_elements_equal/prefix_sums_and_binary_search.rs

@@ -0,0 +1,54 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl Solution {
+    pub fn min_operations(nums: Vec<i32>, queries: Vec<i32>) -> Vec<i64> {
+        let mut nums = nums.into_iter().map(|num| num as u32).collect::<Vec<_>>();
+        let n = nums.len() as u64;
+
+        nums.sort_unstable();
+
+        let mut sum = 0;
+
+        let sums = nums
+            .iter()
+            .map(|&num| {
+                sum += u64::from(num);
+
+                sum
+            })
+            .collect::<Box<_>>();
+
+        queries
+            .into_iter()
+            .map(|query| {
+                let query = u64::from(query as u32);
+                let index = nums.partition_point(|&num| num < query as u32);
+                let left_sum = sums.get(index.wrapping_sub(1)).copied().unwrap_or(0);
+                let index = index as u64;
+                let left_operations = query * index - left_sum;
+                let right_sum = sum - left_sum;
+                let right_operations = right_sum - query * (n - index);
+
+                (left_operations + right_operations) as _
+            })
+            .collect()
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn min_operations(nums: Vec<i32>, queries: Vec<i32>) -> Vec<i64> {
+        Self::min_operations(nums, queries)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}