1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

pub struct ReRooting<T, Identity, Merge, AddRoot> {
    dp: Vec<Vec<T>>,
    ans: Vec<T>,
    graph: Vec<Vec<usize>>,
    identity: Identity,
    merge: Merge,
    add_root: AddRoot,
}

impl<T, Identity, Merge, AddRoot> ReRooting<T, Identity, Merge, AddRoot>
where
    T: Clone,
    Identity: Fn() -> T,
    Merge: Fn(T, T) -> T,
    AddRoot: Fn(T) -> T,
{
    pub fn new(n: usize, identity: Identity, merge: Merge, add_root: AddRoot) -> Self {
        Self {
            dp: vec![vec![]; n],
            ans: vec![identity(); n],
            graph: vec![vec![]; n],
            identity,
            merge,
            add_root,
        }
    }
    pub fn add_edge(&mut self, a: usize, b: usize) {
        self.graph[a].push(b);
    }
    pub fn build(&mut self) {
        self.dfs(0, 0);
        self.dfs2(0, 0, (self.identity)());
    }

    fn dfs(&mut self, v: usize, p: usize) -> T {
        let mut sum = (self.identity)();
        let deg = self.graph[v].len();
        self.dp[v] = vec![(self.identity)(); deg];
        let next = self.graph[v].clone();
        for (i, next) in next.into_iter().enumerate() {
            if next == p {
                continue;
            }
            let t = self.dfs(next, v);
            self.dp[v][i] = t.clone();
            sum = (self.merge)(sum, t);
        }
        (self.add_root)(sum)
    }
    fn dfs2(&mut self, v: usize, p: usize, dp_p: T) {
        for (i, &next) in self.graph[v].iter().enumerate() {
            if next == p {
                self.dp[v][i] = dp_p.clone();
            }
        }

        let deg = self.graph[v].len();
        let mut dp_l = vec![(self.identity)(); deg + 1];
        let mut dp_r = vec![(self.identity)(); deg + 1];
        for i in 0..deg {
            dp_l[i + 1] = (self.merge)(dp_l[i].clone(), self.dp[v][i].clone());
        }
        for i in (0..deg).rev() {
            dp_r[i] = (self.merge)(dp_r[i + 1].clone(), self.dp[v][i].clone());
        }

        self.ans[v] = (self.add_root)(dp_l[deg].clone());

        let next = self.graph[v].clone();
        for (i, next) in next.into_iter().enumerate() {
            if next == p {
                continue;
            }
            self.dfs2(
                next,
                v,
                (self.add_root)((self.merge)(dp_l[i].clone(), dp_r[i + 1].clone())),
            );
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::graph::re_rooting::ReRooting;

    #[test]
    fn test_re_rooting() {
        fn comb(n: usize, k: usize) -> usize {
            let mut ans = 1;
            for i in 0..k {
                ans *= n - i;
                ans /= i + 1;
            }
            ans
        }
        let merge = |e1: Option<(i64, usize)>, e2: Option<(i64, usize)>| {
            if let (Some((ans1, size1)), Some((ans2, size2))) = (e1, e2) {
                let c = comb(size1 + size2, size1);
                let ans = ans1 * ans2 * (c as i64);
                Some((ans, size1 + size2))
            } else {
                e1.or(e2)
            }
        };
        let add_root =
            |e: Option<(i64, usize)>| e.map(|(ans, size)| (ans, size + 1)).or(Some((1, 1)));

        let n = 8;
        let mut graph = ReRooting::new(n, || None, merge, add_root);
        let edges = vec![(1, 2), (2, 3), (3, 4), (3, 5), (3, 6), (6, 7), (6, 8)];
        for (u, v) in edges {
            let u = u - 1;
            let v = v - 1;
            graph.add_edge(u, v);
            graph.add_edge(v, u);
        }

        graph.build();
        assert_eq!(graph.ans[0].unwrap().0, 40);
        assert_eq!(graph.ans[1].unwrap().0, 280);
        assert_eq!(graph.ans[2].unwrap().0, 840);
        assert_eq!(graph.ans[3].unwrap().0, 120);
        assert_eq!(graph.ans[4].unwrap().0, 120);
        assert_eq!(graph.ans[5].unwrap().0, 504);
        assert_eq!(graph.ans[6].unwrap().0, 72);
        assert_eq!(graph.ans[7].unwrap().0, 72);
    }
}