re_rooting_dp/
lib.rs

1/// 全方位木DP
2///
3/// `fold(p, ch, e)` は親頂点 `p` に子の頂点 `ch` を辺 `e` 含めてマージした結果を返すよう実装する
4///
5/// ```no_run
6/// // 各頂点から最も遠い頂点までの距離を求める例
7///
8/// use re_rooting_dp::re_rooting_dp;
9///
10/// struct E(u64);
11/// #[derive(Clone)]
12/// struct V(u64);
13///
14/// let n: usize = todo!();
15/// let edges: Vec<(usize, usize, E)> = todo!();
16///
17/// let farthest = re_rooting_dp(
18///     n,
19///     &edges,
20///     // new
21///     |_i| {
22///         V(0)
23///     },
24///     // fold
25///     |p, ch, e| {
26///         V(p.0.max(ch.0 + e.0))
27///     }
28/// );
29/// ```
30pub fn re_rooting_dp<E, V, F, G>(n: usize, edges: &[(usize, usize, E)], new: F, fold: G) -> Vec<V>
31where
32    V: Clone,
33    F: Fn(usize) -> V,
34    G: Fn(&V, &V, &E) -> V,
35{
36    if n == 0 {
37        return Vec::new();
38    }
39
40    let (g, pre_order) = {
41        let mut g = vec![vec![]; n];
42        for (u, v, e) in edges {
43            g[*u].push((*v, e));
44            g[*v].push((*u, e));
45        }
46        let mut ord = Vec::with_capacity(n);
47        let mut stack = vec![(0, usize::MAX)];
48        while let Some((i, p)) = stack.pop() {
49            ord.push(i);
50            g[i].retain(|&(j, _)| j != p);
51            for &(j, _) in &g[i] {
52                stack.push((j, i));
53            }
54        }
55        (g, ord)
56    };
57
58    // 部分木に対するDP
59    let dp_sub = {
60        let mut dp_sub = (0..n).map(&new).collect::<Vec<_>>();
61        for &i in pre_order.iter().rev() {
62            for &(j, e) in &g[i] {
63                dp_sub[i] = fold(&dp_sub[i], &dp_sub[j], e);
64            }
65        }
66        dp_sub
67    };
68
69    // 親方向に対するDP
70    let mut dp_p = (0..n).map(&new).collect::<Vec<_>>();
71    for i in pre_order {
72        // 頂点iの子である全ての頂点jについてdp_p[j]を更新する
73        apply(dp_p[i].clone(), &g[i], &fold, &dp_sub, &mut dp_p);
74    }
75
76    dp_p.into_iter()
77        .enumerate()
78        .map(|(i, dp_p)| {
79            g[i].iter()
80                .fold(dp_p, |acc, &(j, e)| fold(&acc, &dp_sub[j], e))
81        })
82        .collect::<Vec<_>>()
83}
84
85fn apply<E, V, G>(acc: V, children: &[(usize, &E)], fold: &G, dp_sub: &Vec<V>, dp_p: &mut Vec<V>)
86where
87    V: Clone,
88    G: Fn(&V, &V, &E) -> V,
89{
90    if children.is_empty() {
91        return;
92    }
93
94    if children.len() == 1 {
95        let (j, e) = children[0];
96        dp_p[j] = fold(&dp_p[j], &acc, e);
97        return;
98    }
99
100    let (left, right) = children.split_at(children.len() / 2);
101    let left_acc = left
102        .iter()
103        .fold(acc.clone(), |acc, &(j, e)| fold(&acc, &dp_sub[j], e));
104    let right_acc = right
105        .iter()
106        .fold(acc, |acc, &(j, e)| fold(&acc, &dp_sub[j], e));
107    apply(left_acc, right, fold, dp_sub, dp_p);
108    apply(right_acc, left, fold, dp_sub, dp_p);
109}
re_rooting_dp/lib.rs

re_rooting_dp/
lib.rs
