diff --git a/algo/re_rooting_dp/Cargo.toml b/algo/re_rooting_dp/Cargo.toml new file mode 100644 index 00000000..64c2898f --- /dev/null +++ b/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 + +[dev-dependencies] +proconio = {version = "0.4.5", features = ["derive"] } diff --git a/algo/re_rooting_dp/examples/tree_path_composite_sum.rs b/algo/re_rooting_dp/examples/tree_path_composite_sum.rs new file mode 100644 index 00000000..c5ab0528 --- /dev/null +++ b/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::>(); + + 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::>() + .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, // 部分木のサイズ +} diff --git a/algo/re_rooting_dp/src/lib.rs b/algo/re_rooting_dp/src/lib.rs new file mode 100644 index 00000000..bbc010c6 --- /dev/null +++ b/algo/re_rooting_dp/src/lib.rs @@ -0,0 +1,109 @@ +/// 全方位木DP +/// +/// `fold(p, ch, e)` は親頂点 `p` に子の頂点 `ch` を辺 `e` 含めてマージした結果を返すよう実装する +/// +/// ```no_run +/// // 各頂点から最も遠い頂点までの距離を求める例 +/// +/// use re_rooting_dp::re_rooting_dp; +/// +/// struct E(u64); +/// #[derive(Clone)] +/// struct V(u64); +/// +/// let n: usize = todo!(); +/// let edges: Vec<(usize, usize, E)> = todo!(); +/// +/// let farthest = re_rooting_dp( +/// n, +/// &edges, +/// // new +/// |_i| { +/// V(0) +/// }, +/// // fold +/// |p, ch, e| { +/// V(p.0.max(ch.0 + e.0)) +/// } +/// ); +/// ``` +pub fn re_rooting_dp(n: usize, edges: &[(usize, usize, E)], new: F, fold: G) -> Vec +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::>(); + 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::>(); + 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::>() +} + +fn apply(acc: V, children: &[(usize, &E)], fold: &G, dp_sub: &Vec, dp_p: &mut Vec) +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); +}