dsalgo/
tree_dfs_abstract.rs

1use crate::tree_edges_to_graph::tree_edges_to_graph;
2
3pub fn tree_dfs<T, F>(
4    tree_edges: &[(usize, usize)],
5    root: usize,
6    default_data: Vec<T>,
7    mut assign: F,
8) -> Vec<T>
9where
10    F: FnMut(&mut Vec<T>, usize, usize),
11{
12    let graph = tree_edges_to_graph(tree_edges);
13
14    let n = graph.len();
15
16    assert!(default_data.len() == n);
17
18    let mut parent = vec![None; n];
19
20    let mut data = default_data;
21
22    let mut stack: Vec<isize> = Vec::new();
23
24    stack.push(root as isize);
25
26    while let Some(u) = stack.pop() {
27        if u < 0 {
28            let u = !u as usize;
29
30            if let Some(p) = parent[u] {
31                assign(&mut data, p, u);
32            }
33
34            continue;
35        }
36
37        stack.push(!u);
38
39        let u = u as usize;
40
41        for &v in graph[u].iter() {
42            if Some(v) == parent[u] {
43                continue;
44            }
45
46            parent[v] = Some(u);
47
48            stack.push(v as isize);
49        }
50    }
51
52    data
53}
54
55#[cfg(test)]
56
57mod tests {
58
59    #[test]
60
61    fn test() {}
62}