From 398ac28cc0932c95716cd6873200bfb701dc27aa Mon Sep 17 00:00:00 2001
From: EFanZh <efanzh@gmail.com>
Date: Thu, 25 Jul 2024 23:30:00 +0800
Subject: [PATCH] Add problem 2192: All Ancestors of a Node in a Directed
 Acyclic Graph

---
 src/lib.rs                                    |  1 +
 .../dfs.rs                                    | 73 +++++++++++++++++++
 .../mod.rs                                    | 53 ++++++++++++++
 3 files changed, 127 insertions(+)
 create mode 100644 src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/dfs.rs
 create mode 100644 src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/mod.rs

diff --git a/src/lib.rs b/src/lib.rs
index 86ff70a3..2e975f47 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1577,6 +1577,7 @@ pub mod problem_2185_counting_words_with_a_given_prefix;
 pub mod problem_2186_minimum_number_of_steps_to_make_two_strings_anagram_ii;
 pub mod problem_2190_most_frequent_number_following_key_in_an_array;
 pub mod problem_2191_sort_the_jumbled_numbers;
+pub mod problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph;
 pub mod problem_2194_cells_in_a_range_on_an_excel_sheet;
 pub mod problem_2195_append_k_integers_with_minimal_sum;
 pub mod problem_2196_create_binary_tree_from_descriptions;
diff --git a/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/dfs.rs b/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/dfs.rs
new file mode 100644
index 00000000..55d058be
--- /dev/null
+++ b/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/dfs.rs
@@ -0,0 +1,73 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::collections::HashSet;
+use std::mem;
+
+impl Solution {
+    fn dfs(states: &mut [Result<Vec<u16>, HashSet<u16>>], node: u16) {
+        let node = usize::from(node);
+        let state = &mut states[node];
+
+        if state.is_ok() {
+            let parents = mem::take(state.as_mut().ok().unwrap());
+            let mut ancestors = parents.iter().copied().collect::<HashSet<_>>();
+
+            for parent in parents {
+                Self::dfs(states, parent);
+                ancestors.extend(states[usize::from(parent)].as_ref().unwrap_err().iter().copied());
+            }
+
+            drop(mem::replace(&mut states[node], Err(ancestors)));
+        }
+    }
+
+    pub fn get_ancestors(n: i32, edges: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
+        let n = n as u32 as usize;
+        let mut states = vec![Vec::new(); n];
+
+        for edge in edges {
+            let [from, to]: [_; 2] = edge.try_into().ok().unwrap();
+
+            states[usize::from(to as u16)].push(from as u16);
+        }
+
+        let mut states = states.into_iter().map(Ok).collect::<Vec<_>>();
+
+        for node in 0..n as _ {
+            Self::dfs(&mut states, node);
+        }
+
+        states
+            .into_iter()
+            .map(|state| {
+                let mut ancestors = state
+                    .unwrap_err()
+                    .into_iter()
+                    .map(|x| u32::from(x) as _)
+                    .collect::<Vec<_>>();
+
+                ancestors.sort_unstable();
+
+                ancestors
+            })
+            .collect()
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn get_ancestors(n: i32, edges: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
+        Self::get_ancestors(n, edges)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}
diff --git a/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/mod.rs b/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/mod.rs
new file mode 100644
index 00000000..cd4e65d3
--- /dev/null
+++ b/src/problem_2192_all_ancestors_of_a_node_in_a_directed_acyclic_graph/mod.rs
@@ -0,0 +1,53 @@
+pub mod dfs;
+
+pub trait Solution {
+    fn get_ancestors(n: i32, edges: Vec<Vec<i32>>) -> Vec<Vec<i32>>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            (
+                (
+                    8,
+                    &[[0, 3], [0, 4], [1, 3], [2, 4], [2, 7], [3, 5], [3, 6], [3, 7], [4, 6]] as &[_],
+                ),
+                &[
+                    &[] as &[_],
+                    &[],
+                    &[],
+                    &[0, 1],
+                    &[0, 2],
+                    &[0, 1, 3],
+                    &[0, 1, 2, 3, 4],
+                    &[0, 1, 2, 3],
+                ] as &[&[_]],
+            ),
+            (
+                (
+                    5,
+                    &[
+                        [0, 1],
+                        [0, 2],
+                        [0, 3],
+                        [0, 4],
+                        [1, 2],
+                        [1, 3],
+                        [1, 4],
+                        [2, 3],
+                        [2, 4],
+                        [3, 4],
+                    ],
+                ),
+                &[&[], &[0], &[0, 1], &[0, 1, 2], &[0, 1, 2, 3]],
+            ),
+        ];
+
+        for ((n, edges), expected) in test_cases {
+            assert_eq!(S::get_ancestors(n, edges.iter().map(Vec::from).collect()), expected);
+        }
+    }
+}