haar_lib/tree/
tree_dp.rs

1//! 木DP
2
3use crate::tree::*;
4
5/// 木DP
6///
7/// # Problems
8/// - <https://atcoder.jp/contests/dp/tasks/dp_p>
9/// - <https://yukicoder.me/problems/no/763>
10pub struct TreeDP<'a, T, U, E> {
11    init: U,
12    up: Box<dyn 'a + Fn(T, &'a E) -> U>,
13    merge: Box<dyn 'a + Fn(U, U) -> U>,
14    apply: Box<dyn 'a + Fn(U, usize) -> T>,
15}
16
17impl<'a, T, U, E> TreeDP<'a, T, U, E>
18where
19    E: TreeEdgeTrait,
20    T: Clone,
21    U: Clone,
22{
23    /// `TreeDP`を構築する。
24    pub fn new(
25        init: U,
26        up: Box<impl 'a + Fn(T, &'a E) -> U>,
27        merge: Box<impl 'a + Fn(U, U) -> U>,
28        apply: Box<impl 'a + Fn(U, usize) -> T>,
29    ) -> Self {
30        Self {
31            init,
32            up,
33            merge,
34            apply,
35        }
36    }
37
38    /// `root`を根にして、`tree`上でDPを実行する。
39    ///
40    /// **Time complexity** $O(n)$
41    pub fn run(&self, tree: &'a Tree<E>, root: usize) -> Vec<T> {
42        let size = tree.len();
43        let mut ret = vec![None; size];
44
45        self.__dfs(tree, root, None, &mut ret);
46
47        ret.into_iter().flatten().collect()
48    }
49
50    fn __dfs(
51        &self,
52        tree: &'a Tree<E>,
53        cur: usize,
54        par: Option<usize>,
55        ret: &mut Vec<Option<T>>,
56    ) -> T {
57        let acc = tree.nodes[cur]
58            .neighbors()
59            .filter(|e| par.is_none_or(|p| p != e.to()))
60            .map(|e| {
61                let a = self.__dfs(tree, e.to(), Some(cur), ret);
62                (self.up)(a, e)
63            })
64            .fold(self.init.clone(), |a, b| (self.merge)(a, b));
65        ret[cur] = Some((self.apply)(acc, cur));
66        ret[cur].clone().unwrap()
67    }
68}
haar_lib/tree/tree_dp.rs

haar_lib/tree/
tree_dp.rs
