diff --git a/src/problem/mod.rs b/src/problem/mod.rs index 7705b41..5c58bd5 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -314,3 +314,5 @@ mod p1547_minimum_cost_to_cut_a_stick; mod p3258_count_substrings_that_satisfy_k_constraint_i; mod p3261_count_substrings_that_satisfy_k_constraint_ii; + +mod p3249_count_the_number_of_good_nodes; diff --git a/src/problem/p3249_count_the_number_of_good_nodes.rs b/src/problem/p3249_count_the_number_of_good_nodes.rs new file mode 100644 index 0000000..1f04eb9 --- /dev/null +++ b/src/problem/p3249_count_the_number_of_good_nodes.rs @@ -0,0 +1,97 @@ +/** + * [3249] Count the Number of Good Nodes + */ +pub struct Solution {} + +// submission codes start here + +impl Solution { + pub fn count_good_nodes(edges: Vec>) -> i32 { + let n = edges.len() + 1; + let mut matrix = vec![vec![]; n]; + + edges + .into_iter() + .map(|edge| (edge[0] as usize, edge[1] as usize)) + .for_each(|(x, y)| { + matrix[x].push(y); + matrix[y].push(x); + }); + + let mut result = 0; + Self::dfs(&matrix, 0, usize::MAX, &mut result); + result + } + + fn dfs(tree: &Vec>, node: usize, parent: usize, result: &mut i32) -> i32 { + if tree[node].len() == 0 { + *result += 1; + return 1; + } + + let children: Vec = tree[node] + .iter() + .filter_map(|&child| { + if child == parent { + None + } else { + Some(Self::dfs(tree, child, node, result)) + } + }) + .collect(); + + if children.iter().all(|x| *x == children[0]) { + *result += 1; + } + + children.into_iter().sum::() + 1 + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_3249() { + assert_eq!( + 7, + Solution::count_good_nodes(vec![ + vec![0, 1], + vec![0, 2], + vec![1, 3], + vec![1, 4], + vec![2, 5], + vec![2, 6] + ]) + ); + + assert_eq!( + 6, + Solution::count_good_nodes(vec![ + vec![0, 1], + vec![1, 2], + vec![2, 3], + vec![3, 4], + vec![0, 5], + vec![1, 6], + vec![2, 7], + vec![3, 8] + ]) + ); + + assert_eq!( + 6, + Solution::count_good_nodes(vec![ + vec![6, 0], + vec![1, 0], + vec![5, 1], + vec![2, 5], + vec![3, 1], + vec![4, 3] + ]) + ); + } +}