rustgym/leetcode/
_863_all_nodes_distance_k_in_binary_tree.rs

1struct Solution;
2use rustgym_util::*;
3use std::collections::HashMap;
4use std::collections::HashSet;
5use std::collections::VecDeque;
6
7impl Solution {
8    fn distance_k(root: TreeLink, p: TreeLink, k: i32) -> Vec<i32> {
9        let k = k as usize;
10        let mut adj: HashMap<i32, Vec<i32>> = HashMap::new();
11        Self::dfs(root, &mut adj);
12        let mut res = vec![];
13        let mut queue: VecDeque<(i32, usize)> = VecDeque::new();
14        let start = p.unwrap().borrow().val;
15        let mut visited: HashSet<i32> = HashSet::new();
16        visited.insert(start);
17        queue.push_back((start, 0));
18        while let Some((u, d)) = queue.pop_front() {
19            if d == k {
20                res.push(u);
21            } else {
22                for &mut v in adj.entry(u).or_default() {
23                    if visited.insert(v) {
24                        queue.push_back((v, d + 1));
25                    }
26                }
27            }
28        }
29        res
30    }
31
32    fn dfs(root: TreeLink, adj: &mut HashMap<i32, Vec<i32>>) {
33        if let Some(node) = root {
34            let mut node = node.borrow_mut();
35            let u = node.val;
36            let left = node.left.take();
37            let right = node.right.take();
38            if left.is_some() {
39                let v = left.as_ref().unwrap().borrow().val;
40                adj.entry(u).or_default().push(v);
41                adj.entry(v).or_default().push(u);
42                Self::dfs(left, adj);
43            }
44            if right.is_some() {
45                let v = right.as_ref().unwrap().borrow().val;
46                adj.entry(u).or_default().push(v);
47                adj.entry(v).or_default().push(u);
48                Self::dfs(right, adj);
49            }
50        }
51    }
52}
53
54#[test]
55fn test() {
56    let root = tree!(
57        3,
58        tree!(5, tree!(6), tree!(2, tree!(7), tree!(4))),
59        tree!(1, tree!(0), tree!(8))
60    );
61    let p = tree!(5);
62    let k = 2;
63    let mut res = vec![7, 4, 1];
64    let mut ans = Solution::distance_k(root, p, k);
65    res.sort_unstable();
66    ans.sort_unstable();
67    assert_eq!(ans, res);
68}
rustgym/leetcode/_863_all_nodes_distance_k_in_binary_tree.rs

rustgym/leetcode/
_863_all_nodes_distance_k_in_binary_tree.rs
