Add problem 3275: K-th Nearest Obstacle Queries

Author: EFanZh

src/lib.rs

@@ -2165,6 +2165,7 @@ pub mod problem_3264_final_array_state_after_k_multiplication_operations_i;
 pub mod problem_3270_find_the_key_of_the_numbers;
 pub mod problem_3271_hash_divided_string;
 pub mod problem_3274_check_if_two_chessboard_squares_have_the_same_color;
+pub mod problem_3275_k_th_nearest_obstacle_queries;
 pub mod problem_3350_adjacent_increasing_subarrays_detection_ii;
 
 #[cfg(test)]

src/problem_3275_k_th_nearest_obstacle_queries/binary_heap.rs

@@ -0,0 +1,58 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::collections::BinaryHeap;
+use std::iter;
+
+impl Solution {
+    pub fn results_array(queries: Vec<Vec<i32>>, k: i32) -> Vec<i32> {
+        let k = k.cast_unsigned() as usize;
+
+        if k <= queries.len() {
+            let mut queries = queries
+                .into_iter()
+                .map(|query| query.into_iter().map(i32::unsigned_abs).sum::<u32>());
+
+            let mut queue = queries.by_ref().take(k).collect::<BinaryHeap<_>>();
+
+            iter::repeat_n(-1, k - 1)
+                .chain(Some(queue.peek().unwrap().cast_signed()))
+                .chain(queries.map(|distance| {
+                    let mut peek_mut = queue.peek_mut().unwrap();
+                    let target = &mut *peek_mut;
+
+                    if distance < *target {
+                        *target = distance;
+                    }
+
+                    drop(peek_mut);
+
+                    queue.peek().unwrap().cast_signed()
+                }))
+                .collect()
+        } else {
+            let n = queries.len();
+
+            drop(queries);
+
+            vec![-1; n]
+        }
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn results_array(queries: Vec<Vec<i32>>, k: i32) -> Vec<i32> {
+        Self::results_array(queries, k)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3275_k_th_nearest_obstacle_queries/mod.rs

@@ -0,0 +1,23 @@
+pub mod binary_heap;
+
+pub trait Solution {
+    fn results_array(queries: Vec<Vec<i32>>, k: i32) -> Vec<i32>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((&[[1, 2], [3, 4], [2, 3], [-3, 0]] as &[_], 2), &[-1, 7, 5, 3] as &[_]),
+            ((&[[5, 5], [4, 4], [3, 3]], 1), &[10, 8, 6]),
+            ((&[[-7, -1]], 3), &[-1]),
+            ((&[[-6, 4], [7, 8], [-2, -1], [1, -9], [-9, 4]], 1), &[10, 10, 3, 3, 3]),
+        ];
+
+        for ((queries, k), expected) in test_cases {
+            assert_eq!(S::results_array(queries.iter().map(Vec::from).collect(), k), expected,);
+        }
+    }
+}