re-rooting dp

Author: ia7ck

algo/re_rooting_dp/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "re_rooting_dp"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+proconio = {version = "0.4.5", features = ["derive"] }

algo/re_rooting_dp/examples/tree_path_composite_sum.rs

@@ -0,0 +1,53 @@
+// problem: https://judge.yosupo.jp/problem/tree_path_composite_sum
+use proconio::input;
+use re_rooting_dp::re_rooting_dp;
+
+fn main() {
+    input! {
+        n: usize,
+        a: [u64; n],
+        edges: [(usize, usize, u64, u64); n - 1],
+    };
+
+    const M: u64 = 998244353;
+
+    let edges = edges
+        .into_iter()
+        .map(|(u, v, b, c)| (u, v, E { b, c }))
+        .collect::<Vec<_>>();
+
+    let ans = re_rooting_dp(
+        n,
+        &edges,
+        |i| V { val: a[i], size: 1 },
+        |p, ch, e| {
+            // Σ_j e.b * P(ch, j) + e.c
+            // = e.c * ch.size + e.b * Σ_j P(ch, j)
+
+            V {
+                val: (p.val + e.c * ch.size % M + e.b * ch.val % M) % M,
+                size: p.size + ch.size,
+            }
+        },
+    );
+
+    let ans = ans
+        .iter()
+        .map(|v| v.val.to_string())
+        .collect::<Vec<_>>()
+        .join(" ");
+    println!("{}", ans);
+}
+
+#[derive(Debug)]
+struct E {
+    b: u64,
+    c: u64,
+}
+
+#[derive(Debug, Clone)]
+struct V {
+    // i: usize,
+    val: u64,  // Σ P(i, j)
+    size: u64, // 部分木のサイズ
+}

algo/re_rooting_dp/src/lib.rs

@@ -0,0 +1,104 @@
+/// 全方位木DP
+///
+/// `fold(p, ch, e)` は親頂点 `p` に子の頂点 `ch` を辺 `e` 含めてマージした結果を返すよう実装する
+///
+/// ```no_run
+/// // 木の直径を求める例
+///
+/// struct E(u64);
+/// #[derive(Clone)]
+/// struct V(u64);
+///
+/// re_rooting_dp(
+///     n,
+///     &edges,
+///     // new
+///     |_i| {
+///         V(0)
+///     },
+///     // fold
+///     |p, ch, e| {
+///         p.0.max(ch.0 + e.0)
+///     }
+/// )
+/// ```
+pub fn re_rooting_dp<E, V, F, G>(n: usize, edges: &[(usize, usize, E)], new: F, fold: G) -> Vec<V>
+where
+    V: Clone,
+    F: Fn(usize) -> V,
+    G: Fn(&V, &V, &E) -> V,
+{
+    if n == 0 {
+        return Vec::new();
+    }
+
+    let (g, pre_order) = {
+        let mut g = vec![vec![]; n];
+        for (u, v, e) in edges {
+            g[*u].push((*v, e));
+            g[*v].push((*u, e));
+        }
+        let mut ord = Vec::with_capacity(n);
+        let mut stack = vec![(0, usize::MAX)];
+        while let Some((i, p)) = stack.pop() {
+            ord.push(i);
+            g[i].retain(|&(j, _)| j != p);
+            for &(j, _) in &g[i] {
+                stack.push((j, i));
+            }
+        }
+        (g, ord)
+    };
+
+    // 部分木に対するDP
+    let dp_sub = {
+        let mut dp_sub = (0..n).map(&new).collect::<Vec<_>>();
+        for &i in pre_order.iter().rev() {
+            for &(j, e) in &g[i] {
+                dp_sub[i] = fold(&dp_sub[i], &dp_sub[j], e);
+            }
+        }
+        dp_sub
+    };
+
+    // 親方向に対するDP
+    let mut dp_p = (0..n).map(&new).collect::<Vec<_>>();
+    for i in pre_order {
+        // 頂点iの子である全ての頂点jについてdp_p[j]を更新する
+        apply(dp_p[i].clone(), &g[i], &fold, &dp_sub, &mut dp_p);
+    }
+
+    dp_p.into_iter()
+        .enumerate()
+        .map(|(i, dp_p)| {
+            g[i].iter()
+                .fold(dp_p, |acc, &(j, e)| fold(&acc, &dp_sub[j], e))
+        })
+        .collect::<Vec<_>>()
+}
+
+fn apply<E, V, G>(acc: V, children: &[(usize, &E)], fold: &G, dp_sub: &Vec<V>, dp_p: &mut Vec<V>)
+where
+    V: Clone,
+    G: Fn(&V, &V, &E) -> V,
+{
+    if children.is_empty() {
+        return;
+    }
+
+    if children.len() == 1 {
+        let (j, e) = children[0];
+        dp_p[j] = fold(&dp_p[j], &acc, e);
+        return;
+    }
+
+    let (left, right) = children.split_at(children.len() / 2);
+    let left_acc = left
+        .iter()
+        .fold(acc.clone(), |acc, &(j, e)| fold(&acc, &dp_sub[j], e));
+    let right_acc = right
+        .iter()
+        .fold(acc, |acc, &(j, e)| fold(&acc, &dp_sub[j], e));
+    apply(left_acc, right, fold, dp_sub, dp_p);
+    apply(right_acc, left, fold, dp_sub, dp_p);
+}