competitive/tree/
rerooting.rs

1//! dynamic programming on all-rooted trees
2
3use crate::algebra::Monoid;
4use crate::graph::{Adjacency, UndirectedSparseGraph};
5
6#[codesnip::entry("ReRooting", include("algebra", "SparseGraph"))]
7/// dynamic programming on all-rooted trees
8///
9/// caluculate all subtrees (hanging on the edge) in specific ordering,
10/// each subtree calculated in the order of merge and rooting
11#[derive(Clone, Debug)]
12pub struct ReRooting<'a, M: Monoid, F: Fn(&M::T, usize, Option<usize>) -> M::T> {
13    graph: &'a UndirectedSparseGraph,
14    /// dp\[v\]: result of v-rooted tree
15    pub dp: Vec<M::T>,
16    /// ep\[e\]: result of e-subtree, if e >= n then reversed-e-subtree
17    pub ep: Vec<M::T>,
18    /// rooting(data, vid, (Optional)eid): add root node(vid), result subtree is edge(eid)
19    rooting: F,
20}
21#[codesnip::entry("ReRooting")]
22impl<'a, M, F> ReRooting<'a, M, F>
23where
24    M: Monoid,
25    F: Fn(&M::T, usize, Option<usize>) -> M::T,
26{
27    pub fn new(graph: &'a UndirectedSparseGraph, rooting: F) -> Self {
28        let dp = vec![M::unit(); graph.vertices_size()];
29        let ep = vec![M::unit(); graph.vertices_size() * 2];
30        let mut self_ = Self {
31            graph,
32            dp,
33            ep,
34            rooting,
35        };
36        self_.rerooting();
37        self_
38    }
39    #[inline]
40    fn eidx(&self, u: usize, a: &Adjacency) -> usize {
41        a.id + self.graph.edges_size() * (u > a.to) as usize
42    }
43    #[inline]
44    fn reidx(&self, u: usize, a: &Adjacency) -> usize {
45        a.id + self.graph.edges_size() * (u < a.to) as usize
46    }
47    #[inline]
48    fn merge(&self, x: &M::T, y: &M::T) -> M::T {
49        M::operate(x, y)
50    }
51    #[inline]
52    fn add_subroot(&self, x: &M::T, vid: usize, eid: usize) -> M::T {
53        (self.rooting)(x, vid, Some(eid))
54    }
55    #[inline]
56    fn add_root(&self, x: &M::T, vid: usize) -> M::T {
57        (self.rooting)(x, vid, None)
58    }
59    fn dfs(&mut self, pa: &Adjacency, p: usize) {
60        let u = pa.to;
61        let pi = self.eidx(p, pa);
62        for a in self.graph.adjacencies(u).filter(|a| a.to != p) {
63            let i = self.eidx(u, a);
64            self.dfs(a, u);
65            self.ep[pi] = self.merge(&self.ep[pi], &self.ep[i]);
66        }
67        self.ep[pi] = self.add_subroot(&self.ep[pi], u, pa.id);
68    }
69    fn efs(&mut self, u: usize, p: usize) {
70        let m = self.graph.adjacencies(u).len();
71        let mut left = vec![M::unit(); m + 1];
72        let mut right = vec![M::unit(); m + 1];
73        for (k, a) in self.graph.adjacencies(u).enumerate() {
74            let i = self.eidx(u, a);
75            left[k + 1] = self.merge(&left[k], &self.ep[i]);
76        }
77        for (k, a) in self.graph.adjacencies(u).enumerate().rev() {
78            let i = self.eidx(u, a);
79            right[k] = self.merge(&right[k + 1], &self.ep[i]);
80        }
81        self.dp[u] = self.add_root(&left[m], u);
82        for (k, a) in self.graph.adjacencies(u).enumerate() {
83            if a.to != p {
84                let i = self.reidx(u, a);
85                self.ep[i] = self.merge(&left[k], &right[k + 1]);
86                self.ep[i] = self.add_subroot(&self.ep[i], u, a.id);
87                self.efs(a.to, u);
88            }
89        }
90    }
91    fn rerooting(&mut self) {
92        for a in self.graph.adjacencies(0) {
93            self.dfs(a, 0);
94        }
95        self.efs(0, usize::MAX);
96    }
97}
competitive/tree/rerooting.rs

competitive/tree/
rerooting.rs
