20240229 Finished

This commit is contained in:
jackfiled 2024-02-29 11:20:16 +08:00
parent 72ac17fde7
commit 08da695ca3
2 changed files with 108 additions and 1 deletions

View File

@ -58,4 +58,5 @@ mod p2476_closest_nodes_queries_in_a_binary_search_tree;
mod p938_range_sum_of_bst; mod p938_range_sum_of_bst;
mod p889_construct_binary_tree_from_preorder_and_postorder_traversal; mod p889_construct_binary_tree_from_preorder_and_postorder_traversal;
mod p2867_count_valid_paths_in_a_tree; mod p2867_count_valid_paths_in_a_tree;
mod p2673_make_costs_of_paths_equal_in_a_binary_tree; mod p2673_make_costs_of_paths_equal_in_a_binary_tree;
mod p2581_count_number_of_possible_root_nodes;

View File

@ -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<Vec<i32>>, guesses: Vec<Vec<i32>>, 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<Vec<usize>>, guess_set: &HashSet<i64>, 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<Vec<usize>>, guess_set: &HashSet<i64>, 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);
}
}