competitive/tree/
tree_dp.rs

1use crate::graph::UndirectedSparseGraph;
2
3#[codesnip::entry("tree_dp", include("SparseGraph"))]
4impl UndirectedSparseGraph {
5    pub fn tree_dp_bottom_up<T, F>(&self, root: usize, dp: &mut [T], mut f: F)
6    where
7        F: FnMut(&mut T, &T),
8    {
9        fn dfs<T, F>(g: &UndirectedSparseGraph, u: usize, p: usize, dp: &mut [T], f: &mut F)
10        where
11            F: FnMut(&mut T, &T),
12        {
13            for a in g.adjacencies(u) {
14                if a.to != p {
15                    dfs(g, a.to, u, dp, f);
16                    assert_ne!(u, a.to);
17                    let ptr = dp.as_mut_ptr();
18                    unsafe { f(&mut *ptr.add(u), &*ptr.add(a.to)) };
19                }
20            }
21        }
22        dfs(self, root, !0, dp, &mut f);
23    }
24    pub fn tree_dp_top_down<T, F>(&self, root: usize, dp: &mut [T], mut f: F)
25    where
26        F: FnMut(&mut T, &T),
27    {
28        fn dfs<T, F>(g: &UndirectedSparseGraph, u: usize, p: usize, dp: &mut [T], f: &mut F)
29        where
30            F: FnMut(&mut T, &T),
31        {
32            for a in g.adjacencies(u) {
33                if a.to != p {
34                    assert_ne!(u, a.to);
35                    let ptr = dp.as_mut_ptr();
36                    unsafe { f(&mut *ptr.add(a.to), &*ptr.add(u)) };
37                    dfs(g, a.to, u, dp, f);
38                }
39            }
40        }
41        dfs(self, root, !0, dp, &mut f);
42    }
43}
competitive/tree/tree_dp.rs

competitive/tree/
tree_dp.rs
