20240227 Finished

This commit is contained in:
jackfiled 2024-02-27 09:21:05 +08:00
parent d55cc4e4cd
commit 2c7a678406
2 changed files with 95 additions and 1 deletions

View File

@ -57,3 +57,4 @@ mod p2583_kth_largest_sum_in_a_binary_tree;
mod p2476_closest_nodes_queries_in_a_binary_search_tree; 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;

View File

@ -0,0 +1,93 @@
/**
* [2867] Count Valid Paths in a Tree
*/
pub struct Solution {}
// submission codes start here
impl Solution {
fn dfs(graph: &Vec<Vec<usize>>, primes: &Vec<bool>, seen: &mut Vec<usize>,
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<Vec<i32>>) -> 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]]));
}
}