Compress tree

Cargo.toml

@@ -203,3 +203,11 @@ path = "examples/helpers/lis_yosupo.rs"
 [[example]]
 name = "lis_pop"
 path = "examples/helpers/lis_pop.rs"
+
+[[example]]
+name = "mono_st"
+path = "examples/monotonic/mono_st.rs"
+
+[[example]]
+name = "hld_aux_tree"
+path = "examples/graphs/hld_aux_tree.rs"

examples/graphs/hld_aux_tree.rs

@@ -0,0 +1,79 @@
+// 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 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);
+    }
+}

examples/monotonic/mono_st.rs

@@ -0,0 +1,87 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/staticrmq
+
+use proconio::input;
+use programming_team_code_rust::data_structures::rmq::RMQ;
+use programming_team_code_rust::monotonic::mono_range::mono_range;
+use programming_team_code_rust::monotonic::mono_st::mono_st;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize,
+        a: [u32; n],
+    }
+
+    let rmq = RMQ::new(
+        &(0..n).map(|i| (i, i)).collect::<Vec<_>>(),
+        |(min_i1, max_i1), (min_i2, max_i2)| {
+            let min_idx = if a[min_i1] < a[min_i2] {
+                min_i1
+            } else {
+                min_i2
+            };
+            let max_idx = if a[max_i1] > a[max_i2] {
+                max_i1
+            } else {
+                max_i2
+            };
+            (min_idx, max_idx)
+        },
+    );
+
+    let compares = vec![
+        |x: u32, y: u32| -> bool { x.le(&y) },
+        |x: u32, y: u32| -> bool { x.lt(&y) },
+        |x: u32, y: u32| -> bool { x.ge(&y) },
+        |x: u32, y: u32| -> bool { x.gt(&y) },
+    ];
+
+    let mut le = Vec::new();
+    let mut ri = Vec::new();
+    for &cmp in &compares {
+        le.push(mono_st(&a, |&x, &y| cmp(x, y)));
+        ri.push(mono_range(le.last().unwrap()));
+    }
+    let le = le;
+    let ri = ri;
+
+    for curr_le in &le {
+        let mut count_index = vec![0; n];
+        for i in 0..n {
+            let mut j = i.wrapping_sub(1);
+            while j != curr_le[i] {
+                count_index[j] += 1;
+                j = curr_le[j];
+            }
+        }
+        let mut j = n.wrapping_sub(1);
+        while j != usize::MAX {
+            count_index[j] += 1;
+            j = curr_le[j];
+        }
+        assert_eq!(count_index, vec![1; n]);
+    }
+
+    for _ in 0..q {
+        input! {
+            l: usize,
+            r: usize,
+        }
+
+        let (min_idx, max_idx) = rmq.query(l..r);
+
+        for i in 0..4 {
+            let idx = if i < 2 { min_idx } else { max_idx };
+            let rmq_res = rmq.query(le[i][idx].wrapping_add(1)..ri[i][idx]);
+            assert_eq!(a[if i < 2 { rmq_res.0 } else { rmq_res.1 }], a[idx]);
+            if le[i][idx] != usize::MAX {
+                assert!(compares[i](a[le[i][idx]], a[idx]));
+            }
+            if ri[i][idx] < n {
+                assert!(!compares[i](a[idx], a[ri[i][idx]]));
+            }
+        }
+
+        println!("{}", a[min_idx]);
+    }
+}

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
@@ -200,4 +201,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(|&a, &b| self.tin[a].cmp(&self.tin[b]));
+        let siz = nodes.len();
+        for i in 1..siz {
+            nodes.push(self.lca(nodes[i - 1], nodes[i]));
+        }
+        nodes.sort_by(|&a, &b| self.tin[a].cmp(&self.tin[b]));
+        nodes.dedup();
+        (mono_st(&nodes, |&x, &y| self.in_sub(x, y)), nodes)
+    }
 }

src/lib.rs

@@ -2,5 +2,6 @@
 pub mod data_structures;
 pub mod graphs;
 pub mod helpers;
+pub mod monotonic;
 pub mod numbers;
 pub mod strings;

src/monotonic/mod.rs

@@ -0,0 +1,3 @@
+//! # Monotonic
+pub mod mono_range;
+pub mod mono_st;

src/monotonic/mono_range.rs

@@ -0,0 +1,30 @@
+//! # Monotonic Range
+
+/// Gets vec ri where ri\[i\] = smallest index such that i < ri\[i\] && !a\[i\].cmp(&a\[ri\[i\]\]), or n
+///
+/// # Example
+/// ```
+/// use programming_team_code_rust::monotonic::mono_st::mono_st;
+/// use programming_team_code_rust::monotonic::mono_range::mono_range;
+///
+/// let a: Vec<u32> = vec![3, 1, 2, 2];
+///
+/// let le = mono_st(&a, |x, y| x.lt(y));
+/// let ri = mono_range(&le);
+/// assert_eq!(ri, [1, 4, 3, 4]);
+/// ```
+///
+/// # Complexity
+/// - Time: O(n)
+/// - Space: O(n)
+pub fn mono_range(le: &[usize]) -> Vec<usize> {
+    let mut ri = vec![le.len(); le.len()];
+    for i in 0..le.len() {
+        let mut j = i.wrapping_sub(1);
+        while j != le[i] {
+            ri[j] = i;
+            j = le[j];
+        }
+    }
+    ri
+}

src/monotonic/mono_st.rs

@@ -0,0 +1,51 @@
+//! # Monotonic Stack
+
+/// Gets vec le where le\[i\] = largest index such that le\[i\] < i && a\[le\[i\]\].cmp(&a\[i\]), or usize::MAX
+///
+/// # Example
+/// ```
+/// use programming_team_code_rust::monotonic::mono_st::mono_st;
+///
+/// let a: Vec<u32> = vec![3, 1, 2, 2];
+/// let n = a.len();
+///
+/// assert_eq!(mono_st(&a, |x, y| x.lt(y)), [usize::MAX, usize::MAX, 1, 1]);
+/// assert_eq!(mono_st(&a, |x, y| x.le(y)), [usize::MAX, usize::MAX, 1, 2]);
+/// assert_eq!(mono_st(&a, |x, y| x.gt(y)), [usize::MAX, 0, 0, 0]);
+/// assert_eq!(mono_st(&a, |x, y| x.ge(y)), [usize::MAX, 0, 0, 2]);
+///
+/// // simulate popping off stack for single index
+/// let le = mono_st(&a, |x, y| x.lt(y));
+/// let mut seen_index = vec![0; n];
+/// for i in 0..n {
+///    let mut j = i.wrapping_sub(1);
+///    while j != le[i] {
+///       let range = le[j].wrapping_add(1)..j; // for all indexes k in range: !cmp(&a[k], &a[j])
+///       seen_index[j] += 1;
+///       j = le[j];
+///    }
+/// }
+///
+/// // clear the stack at the end
+/// let mut j = n.wrapping_sub(1);
+/// while j != usize::MAX {
+///    seen_index[j] += 1;
+///    j = le[j];
+/// }
+///
+/// assert_eq!(seen_index, vec![1; n]);
+/// ```
+///
+/// # Complexity
+/// - Time: O(n)
+/// - Space: O(n)
+pub fn mono_st<T: Ord, F: Fn(&T, &T) -> bool>(a: &[T], cmp: F) -> Vec<usize> {
+    let mut le = vec![0; a.len()];
+    for i in 0..a.len() {
+        le[i] = i.wrapping_sub(1);
+        while le[i] != usize::MAX && !cmp(&a[le[i]], &a[i]) {
+            le[i] = le[le[i]];
+        }
+    }
+    le
+}