Skip to content

Commit

Permalink
Add problem 2192: All Ancestors of a Node in a Directed Acyclic Graph
Browse files Browse the repository at this point in the history
  • Loading branch information
EFanZh committed Jul 25, 2024
1 parent 1146bf3 commit 398ac28
Show file tree
Hide file tree
Showing 3 changed files with 127 additions and 0 deletions.
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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>();
}
}
Original file line number Diff line number Diff line change
@@ -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);
}
}
}

0 comments on commit 398ac28

Please sign in to comment.