competitive/tree/
depth.rs

1use crate::algebra::Monoid;
2use crate::graph::UndirectedSparseGraph;
3
4#[codesnip::entry("tree_depth", include("SparseGraph"))]
5impl UndirectedSparseGraph {
6    fn depth_dfs(&self, u: usize, p: usize, d: u64, depth: &mut Vec<u64>) {
7        depth[u] = d;
8        for a in self.adjacencies(u).filter(|a| a.to != p) {
9            self.depth_dfs(a.to, u, d + 1, depth);
10        }
11    }
12    pub fn tree_depth(&self, root: usize) -> Vec<u64> {
13        let mut depth = vec![0; self.vertices_size()];
14        self.depth_dfs(root, self.vertices_size(), 0, &mut depth);
15        depth
16    }
17}
18
19#[codesnip::entry("weighted_tree_depth", include("algebra", "SparseGraph"))]
20impl UndirectedSparseGraph {
21    fn weighted_depth_dfs<M, F>(
22        &self,
23        u: usize,
24        p: usize,
25        d: M::T,
26        depth: &mut Vec<M::T>,
27        weight: &F,
28    ) where
29        M: Monoid,
30        F: Fn(usize) -> M::T,
31    {
32        for a in self.adjacencies(u).filter(|a| a.to != p) {
33            let nd = M::operate(&d, &weight(a.id));
34            self.weighted_depth_dfs::<M, _>(a.to, u, nd, depth, weight);
35        }
36        depth[u] = d;
37    }
38    pub fn weighted_tree_depth<M: Monoid, F: Fn(usize) -> M::T>(
39        &self,
40        root: usize,
41        weight: F,
42    ) -> Vec<M::T> {
43        let mut depth = vec![M::unit(); self.vertices_size()];
44        self.weighted_depth_dfs::<M, _>(root, usize::MAX, M::unit(), &mut depth, &weight);
45        depth
46    }
47}
48
49#[codesnip::entry("tree_size", include("SparseGraph"))]
50impl UndirectedSparseGraph {
51    fn size_dfs(&self, u: usize, p: usize, size: &mut Vec<u64>) {
52        size[u] = 1;
53        for a in self.adjacencies(u).filter(|a| a.to != p) {
54            self.size_dfs(a.to, u, size);
55            size[u] += size[a.to];
56        }
57    }
58    pub fn tree_size(&self, root: usize) -> Vec<u64> {
59        let mut size = vec![0; self.vertices_size()];
60        self.size_dfs(root, usize::MAX, &mut size);
61        size
62    }
63}
competitive/tree/depth.rs

competitive/tree/
depth.rs
