diff --git a/src/problem/mod.rs b/src/problem/mod.rs index bdcf592..c81c3f4 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -56,4 +56,5 @@ mod p235_lowest_common_ancestor_of_a_binary_search_tree; mod p2583_kth_largest_sum_in_a_binary_tree; 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; \ No newline at end of file +mod p889_construct_binary_tree_from_preorder_and_postorder_traversal; +mod p2867_count_valid_paths_in_a_tree; \ No newline at end of file diff --git a/src/problem/p2867_count_valid_paths_in_a_tree.rs b/src/problem/p2867_count_valid_paths_in_a_tree.rs new file mode 100644 index 0000000..a4b51be --- /dev/null +++ b/src/problem/p2867_count_valid_paths_in_a_tree.rs @@ -0,0 +1,93 @@ +/** + * [2867] Count Valid Paths in a Tree + */ +pub struct Solution {} + + +// submission codes start here + +impl Solution { + fn dfs(graph: &Vec>, primes: &Vec, seen: &mut Vec, + i: usize, pre: usize) { + seen.push(i); + + for next in &graph[i] { + let next = *next; + if next != pre && !primes[next] { + Solution::dfs(graph, primes, seen, next, i); + } + } + } + + pub fn count_paths(n: i32, edges: Vec>) -> i64 { + let n = n as usize; + let mut primes = vec![true;n + 1]; + primes[1] = false; + for i in 2..=n { + if primes[i] { + if i * i > n { + continue; + } + + for j in (i * i..=n).step_by(i) { + primes[j] = false; + } + } + } + primes = dbg!(primes); + + let mut graph = vec![vec![];n + 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 count = vec![0; n + 1]; + let mut result = 0; + + for root in 1..=n { + if !primes[root] { + continue; + } + + let mut current = 0; + for i in &graph[root] { + let i = *i; + if (primes[i]) { + continue; + } + + if count[i] == 0 { + let mut seen = Vec::new(); + Solution::dfs(&graph, &primes, &mut seen, i, 0); + for k in &seen { + count[*k] = seen.len() as i64; + } + } + + result += count[i] * current; + current += count[i]; + } + + result += current; + } + + result + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_2867() { + assert_eq!(6, Solution::count_paths(6, + vec![vec![1,2],vec![1,3],vec![2,4],vec![3,5],vec![3,6]])); + } +}