diff --git a/Cargo.toml b/Cargo.toml index e0b278d..fc3ee24 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -212,6 +212,10 @@ path = "examples/data_structures/range_container_handmade.rs" name = "mono_st" path = "examples/monotonic/mono_st.rs" +[[example]] +name = "hld_aux_tree" +path = "examples/graphs/hld_aux_tree.rs" + [[example]] name = "count_rects" path = "examples/monotonic/count_rects.rs" diff --git a/examples/graphs/hld_aux_tree.rs b/examples/graphs/hld_aux_tree.rs new file mode 100644 index 0000000..deed9fa --- /dev/null +++ b/examples/graphs/hld_aux_tree.rs @@ -0,0 +1,80 @@ +// verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/problems/GRL_5_B + +use proconio::input; +use programming_team_code_rust::graphs::hld::HLD; +use std::collections::VecDeque; + +fn main() { + input! { + n: usize, + edges: [(usize, usize, u32); n - 1], + } + + let mut adj_weighted = vec![vec![]; n]; + let mut adj = vec![vec![]; n]; + for &(u, v, w) in &edges { + adj[u].push(v); + adj[v].push(u); + adj_weighted[u].push((v, w)); + adj_weighted[v].push((u, w)); + } + let adj_weighted = adj_weighted; + + let mut dist = vec![0; n]; + { + fn dfs(u: usize, p: Option, adj_weighted: &[Vec<(usize, u32)>], dist: &mut [u32]) { + for &(v, w) in &adj_weighted[u] { + if Some(v) == p { + continue; + } + dist[v] = w + dist[u]; + dfs(v, Some(u), adj_weighted, dist); + } + } + dfs(0, None, &adj_weighted, &mut dist); + } + let dist = dist; + + let hld = HLD::new(&mut adj, true); + + let weighted_dist = |u: usize, v: usize| -> u32 { + let lc = hld.lca(u, v); + dist[u] + dist[v] - 2 * dist[lc] + }; + + let mut diam_u = 0; + for i in 1..n { + if weighted_dist(0, i) > weighted_dist(0, diam_u) { + diam_u = i; + } + } + let mut diam_v = 0; + for i in 1..n { + if weighted_dist(diam_u, i) > weighted_dist(diam_u, diam_v) { + diam_v = i; + } + } + + for u in 0..n { + let (par, to_node) = hld.aux_tree(vec![diam_u, diam_v, u]); + let mut aux_adj = vec![vec![]; par.len()]; + for i in 1..par.len() { + let edge_w = dist[to_node[i]] - dist[to_node[par[i]]]; + aux_adj[i].push((par[i], edge_w)); + aux_adj[par[i]].push((i, edge_w)); + } + let mut q = VecDeque::new(); + q.push_back((to_node.iter().position(|&x| x == u).unwrap(), None, 0)); + let mut res = 0; + while let Some((node, parent, curr_dist)) = q.pop_front() { + res = res.max(curr_dist); + for &(v, w) in &aux_adj[node] { + if Some(v) == parent { + continue; + } + q.push_back((v, Some(node), curr_dist + w)); + } + } + println!("{}", res); + } +} diff --git a/src/graphs/hld.rs b/src/graphs/hld.rs index c315a98..eb93a06 100644 --- a/src/graphs/hld.rs +++ b/src/graphs/hld.rs @@ -1,6 +1,7 @@ //! # Heavy Light Decomposition use crate::graphs::dfs_order::{get_dfs_postorder, get_dfs_preorder}; +use crate::monotonic::mono_st::mono_st; use std::ops::Range; /// # Example @@ -201,4 +202,42 @@ impl HLD { None } } + + /// # Auxiliary Tree + /// + /// - see + /// + /// # Example + /// ``` + /// use programming_team_code_rust::graphs::hld::HLD; + /// + /// let n = 5; + /// let mut adj = vec![vec![]; n]; + /// for (u, v) in [(0,1), (1,2), (2,3), (2,4)] { + /// adj[u].push(v); + /// adj[v].push(u); + /// } + /// + /// let hld = HLD::new(&mut adj, false); + /// + /// let (par, to_node) = hld.aux_tree(vec![0, 3, 4]); + /// // 0, 1, .., par.len()-1 is a topological/dfs order of aux tree + /// assert_eq!(par, [usize::MAX, 0, 1, 1]); + /// assert_eq!(to_node, [0, 2, 3, 4]); + /// ``` + /// + /// # Complexity + /// - k = nodes.len() + /// - Time: O((k log k) + (k log n)) + /// - Space: O(k) + pub fn aux_tree(&self, mut nodes: Vec) -> (Vec, Vec) { + nodes.sort_by(|&u, &v| self.tin[u].cmp(&self.tin[v])); + let siz = nodes.len(); + for i in 1..siz { + nodes.push(self.lca(nodes[i - 1], nodes[i])); + } + nodes.sort_by(|&u, &v| self.tin[u].cmp(&self.tin[v])); + nodes.dedup(); + (mono_st(&nodes, |&u, &v| self.in_sub(u, v)), nodes) + } }