From 02fc1ae6fd68e0fa124cd73e207b0012c55a74e9 Mon Sep 17 00:00:00 2001 From: ia7ck <23146842+ia7ck@users.noreply.github.com> Date: Fri, 1 Aug 2025 00:17:38 +0900 Subject: [PATCH 1/4] re-rooting dp --- algo/re_rooting_dp/Cargo.toml | 9 ++ .../examples/tree_path_composite_sum.rs | 53 +++++++++ algo/re_rooting_dp/src/lib.rs | 104 ++++++++++++++++++ 3 files changed, 166 insertions(+) create mode 100644 algo/re_rooting_dp/Cargo.toml create mode 100644 algo/re_rooting_dp/examples/tree_path_composite_sum.rs create mode 100644 algo/re_rooting_dp/src/lib.rs diff --git a/algo/re_rooting_dp/Cargo.toml b/algo/re_rooting_dp/Cargo.toml new file mode 100644 index 00000000..0bbe7521 --- /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 + +[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..81a35c56 --- /dev/null +++ b/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(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); +} From 9a81a59aae8cb44810a60252c93ab34203a4e20b Mon Sep 17 00:00:00 2001 From: ia7ck <23146842+ia7ck@users.noreply.github.com> Date: Fri, 1 Aug 2025 00:20:44 +0900 Subject: [PATCH 2/4] dev-dependencies --- algo/re_rooting_dp/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/algo/re_rooting_dp/Cargo.toml b/algo/re_rooting_dp/Cargo.toml index 0bbe7521..64c2898f 100644 --- a/algo/re_rooting_dp/Cargo.toml +++ b/algo/re_rooting_dp/Cargo.toml @@ -5,5 +5,5 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html -[dependencies] +[dev-dependencies] proconio = {version = "0.4.5", features = ["derive"] } From cf672521d126d83760621a77aeabc241fec6f031 Mon Sep 17 00:00:00 2001 From: ia7ck <23146842+ia7ck@users.noreply.github.com> Date: Fri, 1 Aug 2025 00:26:15 +0900 Subject: [PATCH 3/4] fix doc test --- algo/re_rooting_dp/src/lib.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/algo/re_rooting_dp/src/lib.rs b/algo/re_rooting_dp/src/lib.rs index 81a35c56..8a47922f 100644 --- a/algo/re_rooting_dp/src/lib.rs +++ b/algo/re_rooting_dp/src/lib.rs @@ -3,13 +3,18 @@ /// `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!(); /// -/// re_rooting_dp( +/// let farthest = re_rooting_dp( /// n, /// &edges, /// // new @@ -18,9 +23,9 @@ /// }, /// // fold /// |p, ch, e| { -/// p.0.max(ch.0 + e.0) +/// 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 From b424b5c138f51088afd89a451b53edf01a62632a Mon Sep 17 00:00:00 2001 From: ia7ck <23146842+ia7ck@users.noreply.github.com> Date: Fri, 1 Aug 2025 00:26:57 +0900 Subject: [PATCH 4/4] fmt --- algo/re_rooting_dp/src/lib.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/algo/re_rooting_dp/src/lib.rs b/algo/re_rooting_dp/src/lib.rs index 8a47922f..bbc010c6 100644 --- a/algo/re_rooting_dp/src/lib.rs +++ b/algo/re_rooting_dp/src/lib.rs @@ -4,13 +4,13 @@ /// /// ```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!(); ///