aboutsummaryrefslogtreecommitdiff
path: root/other/num_nodes_sub_tree.rs
diff options
context:
space:
mode:
Diffstat (limited to 'other/num_nodes_sub_tree.rs')
-rw-r--r--other/num_nodes_sub_tree.rs39
1 files changed, 39 insertions, 0 deletions
diff --git a/other/num_nodes_sub_tree.rs b/other/num_nodes_sub_tree.rs
new file mode 100644
index 0000000..ca52f7b
--- /dev/null
+++ b/other/num_nodes_sub_tree.rs
@@ -0,0 +1,39 @@
+use std::collections::{HashMap, HashSet};
+struct Solution;
+
+impl Solution {
+ pub fn count_sub_trees(n: i32, edges: Vec<Vec<i32>>, labels: String) -> Vec<i32> {
+ let mut map = HashMap::<usize, Vec<usize>>::new();
+ for e in &edges {
+ map.entry(e[0] as usize).or_default().push(e[1] as usize);
+ map.entry(e[1] as usize).or_default().push(e[0] as usize);
+ }
+ let mut visited = HashSet::<usize>::new();
+ let mut ans = vec![0; n as usize];
+ Self::dfs(n, &mut map, &labels, 0, &mut visited, &mut ans);
+ ans
+ }
+ fn dfs(n: i32, map: &mut HashMap<usize, Vec<usize>>, labels: &str,
+ x: usize, visited: &mut HashSet<usize>, ans: &mut Vec<i32>) -> [i32; 26] {
+ let mut vec = [0;26];
+ if visited.insert(x) {
+ for e in map.entry(x).or_default().clone() {
+ let v = Self::dfs(n, map, labels, e, visited, ans);
+ for i in 0..26 {
+ vec[i] += v[i]
+ }
+ }
+ let ch = labels.as_bytes()[x];
+ vec[(ch - b'a') as usize] += 1;
+ ans[x] = vec[(ch - b'a') as usize];
+ }
+ vec
+ }
+}
+
+fn main() {
+ let edges = vec![vec![0,1],vec![0,2],vec![1,4],
+ vec![1,5],vec![2,3],vec![2,6]];
+ print!("{:?}", Solution::count_sub_trees(7, edges, String::from("abaedcd")));
+ //expect: [2,1,1,1,1,1,1]
+}