Compress tree (#83)

* saving progress

* saving progress

* Update mono_st.rs

* saving progress

* saving progress

* saving progress

* another assert

* remove

* add test

* asdf

* finish

* better

* fix

* uncomment

* Update mono_st.rs

* Update mono_range.rs

* Update mono_st.rs

* fix

* change style

* nit

* change

* nits

* different style

* golf

* consistency with c++ PTC

* finish doc for mono st

* simplify test

* add

* add docs for mono range

* fix typo

* nit

* fix

* nit to docs

* now ACs

* virt tree

* finish docs

* finished test

* make dist const

* consistent var names

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>
Co-authored-by: Cameron Custer <73217097+cameroncuster@users.noreply.github.com>

Author: lrvideckis

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"

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<usize>, 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);
+    }
+}

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 <https://github.com/kth-competitive-programming/kactl/blob/main/content/graph/CompressTree.h>
+    ///
+    /// # 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<usize>) -> (Vec<usize>, Vec<usize>) {
+        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)
+    }
 }