Add problem 1697: Checking Existence of Edge Length Limited Paths

src/lib.rs

@@ -1269,6 +1269,7 @@ pub mod problem_1691_maximum_height_by_stacking_cuboids;
 pub mod problem_1694_reformat_phone_number;
 pub mod problem_1695_maximum_erasure_value;
 pub mod problem_1696_jump_game_vi;
+pub mod problem_1697_checking_existence_of_edge_length_limited_paths;
 pub mod problem_1700_number_of_students_unable_to_eat_lunch;
 pub mod problem_1701_average_waiting_time;
 pub mod problem_1702_maximum_binary_string_after_change;

src/problem_1697_checking_existence_of_edge_length_limited_paths/mod.rs

@@ -0,0 +1,107 @@
+pub mod union_find;
+
+pub trait Solution {
+    fn distance_limited_paths_exist(n: i32, edge_list: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<bool>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            (
+                (
+                    3,
+                    &[[0, 1, 2], [1, 2, 4], [2, 0, 8], [1, 0, 16]] as &[_],
+                    &[[0, 1, 2], [0, 2, 5]] as &[_],
+                ),
+                &[false, true] as &[_],
+            ),
+            (
+                (
+                    5,
+                    &[[0, 1, 10], [1, 2, 5], [2, 3, 9], [3, 4, 13]],
+                    &[[0, 4, 14], [1, 4, 13]],
+                ),
+                &[true, false],
+            ),
+            (
+                (
+                    13,
+                    &[
+                        [9, 1, 53],
+                        [3, 2, 66],
+                        [12, 5, 99],
+                        [9, 7, 26],
+                        [1, 4, 78],
+                        [11, 1, 62],
+                        [3, 10, 50],
+                        [12, 1, 71],
+                        [12, 6, 63],
+                        [1, 10, 63],
+                        [9, 10, 88],
+                        [9, 11, 59],
+                        [1, 4, 37],
+                        [4, 2, 63],
+                        [0, 2, 26],
+                        [6, 12, 98],
+                        [9, 11, 99],
+                        [4, 5, 40],
+                        [2, 8, 25],
+                        [4, 2, 35],
+                        [8, 10, 9],
+                        [11, 9, 25],
+                        [10, 11, 11],
+                        [7, 6, 89],
+                        [2, 4, 99],
+                        [10, 4, 63],
+                    ],
+                    &[
+                        [9, 7, 65],
+                        [9, 6, 1],
+                        [4, 5, 34],
+                        [10, 8, 43],
+                        [3, 7, 76],
+                        [4, 2, 15],
+                        [7, 6, 52],
+                        [2, 0, 50],
+                        [7, 6, 62],
+                        [1, 0, 81],
+                        [4, 5, 35],
+                        [0, 11, 86],
+                        [12, 5, 50],
+                        [11, 2, 2],
+                        [9, 5, 6],
+                        [12, 0, 95],
+                        [10, 6, 9],
+                        [9, 4, 73],
+                        [6, 10, 48],
+                        [12, 0, 91],
+                        [9, 10, 58],
+                        [9, 8, 73],
+                        [2, 3, 44],
+                        [7, 11, 83],
+                        [5, 3, 14],
+                        [6, 2, 33],
+                    ],
+                ),
+                &[
+                    true, false, false, true, true, false, false, true, false, true, false, true, false, false, false,
+                    true, false, true, false, true, true, true, false, true, false, false,
+                ],
+            ),
+        ];
+
+        for ((n, edge_list, queries), expected) in test_cases {
+            assert_eq!(
+                S::distance_limited_paths_exist(
+                    n,
+                    edge_list.iter().map(Vec::from).collect(),
+                    queries.iter().map(Vec::from).collect()
+                ),
+                expected,
+            );
+        }
+    }
+}

src/problem_1697_checking_existence_of_edge_length_limited_paths/union_find.rs

@@ -0,0 +1,108 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::convert::TryInto;
+
+impl Solution {
+    fn find_root_helper(union_find: &mut [(u32, u32)], node: u32) -> Option<(u32, u32)> {
+        let node = node as usize;
+
+        union_find.get(node).copied().map(|(parent_plus_1, rank)| {
+            let root = if let Some((root, _)) = Self::find_root_helper(union_find, parent_plus_1.wrapping_sub(1)) {
+                union_find[node].0 = root + 1;
+
+                root
+            } else {
+                node as _
+            };
+
+            (root, rank)
+        })
+    }
+
+    fn find_root(union_find: &mut [(u32, u32)], node: u32) -> (u32, u32) {
+        Self::find_root_helper(union_find, node).unwrap()
+    }
+
+    fn union(union_find: &mut [(u32, u32)], x: u32, y: u32) {
+        let left = Self::find_root(union_find, x);
+        let right = Self::find_root(union_find, y);
+
+        if left.0 != right.0 {
+            let (child, parent) = if left.1 < right.1 {
+                (left.0, right.0)
+            } else {
+                if left.1 == right.1 {
+                    union_find[left.0 as usize].1 += 1;
+                }
+
+                (right.0, left.0)
+            };
+
+            union_find[child as usize].0 = parent + 1;
+        }
+    }
+
+    pub fn distance_limited_paths_exist(n: i32, edge_list: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<bool> {
+        let mut edge_list = edge_list
+            .into_iter()
+            .map(|edge| {
+                let [from, to, distance]: [_; 3] = edge.try_into().ok().unwrap();
+
+                (distance as u32, from as u32, to as u32)
+            })
+            .collect::<Box<_>>();
+
+        edge_list.sort_unstable_by_key(|&(distance, ..)| distance);
+
+        let mut result = vec![false; queries.len()];
+
+        let mut queue = result
+            .iter_mut()
+            .zip(queries)
+            .map(|(slot, query)| {
+                let [from, to, limit]: [_; 3] = query.try_into().ok().unwrap();
+
+                (limit as u32, from as u32, to as u32, slot)
+            })
+            .collect::<Vec<_>>();
+
+        queue.sort_unstable_by_key(|&(limit, ..)| limit);
+
+        let mut iter = edge_list.iter();
+        let mut union_find = vec![(0, 0); n as u32 as usize].into_boxed_slice();
+
+        for (limit, p, q, slot) in queue {
+            while let Some(item) = iter.as_slice().first() {
+                if item.0 < limit {
+                    Self::union(&mut union_find, item.1, item.2);
+
+                    iter.next();
+                } else {
+                    break;
+                }
+            }
+
+            *slot = Self::find_root(&mut union_find, p).0 == Self::find_root(&mut union_find, q).0;
+        }
+
+        result
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn distance_limited_paths_exist(n: i32, edge_list: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<bool> {
+        Self::distance_limited_paths_exist(n, edge_list, queries)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}