diff --git a/src/problem/mod.rs b/src/problem/mod.rs index 55c91fa..02a6878 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -58,4 +58,5 @@ mod p2476_closest_nodes_queries_in_a_binary_search_tree; mod p938_range_sum_of_bst; mod p889_construct_binary_tree_from_preorder_and_postorder_traversal; mod p2867_count_valid_paths_in_a_tree; -mod p2673_make_costs_of_paths_equal_in_a_binary_tree; \ No newline at end of file +mod p2673_make_costs_of_paths_equal_in_a_binary_tree; +mod p2581_count_number_of_possible_root_nodes; \ No newline at end of file diff --git a/src/problem/p2581_count_number_of_possible_root_nodes.rs b/src/problem/p2581_count_number_of_possible_root_nodes.rs new file mode 100644 index 0000000..5cc1659 --- /dev/null +++ b/src/problem/p2581_count_number_of_possible_root_nodes.rs @@ -0,0 +1,106 @@ +use std::usize; + +/** + * [2581] Count Number of Possible Root Nodes + */ +pub struct Solution {} + + +// submission codes start here + +use std::collections::HashSet; + +impl Solution { + pub fn root_count(edges: Vec>, guesses: Vec>, k: i32) -> i32 { + let mut graph = vec![vec![];edges.len() + 1]; + + for edge in edges { + let x = edge[0] as usize; + let y = edge[1] as usize; + + graph[x].push(y); + graph[y].push(x); + } + + let mut guess_set = HashSet::with_capacity(guesses.len()); + for guess in guesses { + guess_set.insert(Solution::hash( + guess[0] as usize, + guess[1] as usize)); + } + + let mut count = 0; + Solution::dfs(&graph, &guess_set, 0, usize::MAX, &mut count); + dbg!(count); + + let mut result = 0; + Solution::tree_dp(&graph, &guess_set, &k, 0, + usize::MAX, count, &mut result); + + result + } + + fn hash(x: usize, y: usize) -> i64 { + (x as i64) * 1000000 + (y as i64) + } + + fn dfs(graph: &Vec>, guess_set: &HashSet, now: usize, pre: usize, count: &mut i32) { + for next in &graph[now] { + let next = *next; + if next == pre { + continue; + } + + if guess_set.contains(&Solution::hash(now, next)) { + *count += 1; + } + + Solution::dfs(graph, guess_set, next, now, count); + } + } + + fn tree_dp(graph: &Vec>, guess_set: &HashSet, k: &i32, + now: usize, pre: usize, count: i32, result: &mut i32) { + if count >= *k { + *result += 1; + } + + for next in &graph[now] { + let next = *next; + if next == pre { + continue; + } + + let mut count = count; + if guess_set.contains(&Solution::hash(now, next)) { + count -= 1; + } + + if guess_set.contains(&Solution::hash(next, now)) { + count += 1; + } + + Solution::tree_dp(graph, guess_set, k, next, now, count, result); + } + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_2581() { + assert_eq!( + Solution::root_count( + vec![vec![0,1],vec![1,2],vec![1,3],vec![4,2]], + vec![vec![1,3],vec![0,1],vec![1,0],vec![2,4]], 3), 3); + + assert_eq!( + Solution::root_count( + vec![vec![0,1],vec![2,0],vec![0,3],vec![4,2],vec![3,5],vec![6,0],vec![1,7],vec![2,8],vec![2,9],vec![4,10],vec![9,11],vec![3,12],vec![13,8],vec![14,9],vec![15,9],vec![10,16]], + vec![vec![8,2],vec![12,3],vec![0,1],vec![16,10]], 2), 4); + } +}