The implementation of querying Lowest Common Ancestor.

lca.rs

use std::num::NonZeroUsize;

/// Calculates the lowest common ancestor on given graph.
#[derive(Debug)]
pub struct LowestCommonAncestor {
    parents_by_step: Vec<Vec<Option<NonZeroUsize>>>,
    distances: Vec<usize>,
}

impl LowestCommonAncestor {
    /// Creates a calculator for the lowest common ancestors.
    pub fn new(graph: &[Vec<NonZeroUsize>], root: NonZeroUsize) -> Self {
        let vertices = graph.len();
        let max_step = vertices.next_power_of_two();
        let mut parents_by_step = vec![vec![None; vertices + 1]; max_step];
        let mut distances = vec![0; vertices + 1];

        fn dfs(
            parents_by_step: &mut Vec<Vec<Option<NonZeroUsize>>>,
            distances: &mut Vec<usize>,
            graph: &[Vec<NonZeroUsize>],
            vertex: NonZeroUsize,
            parent: Option<NonZeroUsize>,
            distance: usize,
        ) {
            let vertex_raw = vertex.get();
            parents_by_step[0][vertex_raw] = parent;
            distances[vertex_raw] = distance;
            for &adj in &graph[vertex_raw] {
                if Some(adj) != parent {
                    dfs(
                        parents_by_step,
                        distances,
                        graph,
                        adj,
                        Some(vertex),
                        distance + 1,
                    );
                }
            }
        }
        dfs(&mut parents_by_step, &mut distances, graph, root, None, 0);

        for step in 0..max_step - 1 {
            for vertex in 0..vertices {
                parents_by_step[step + 1][vertex] = parents_by_step[step][vertex]
                    .as_ref()
                    .and_then(|&parent| parents_by_step[step][parent.get()]);
            }
        }
        Self {
            parents_by_step,
            distances,
        }
    }

    /// Queries the lowest common ancestor between the vertex `u` and `v`.
    pub fn query(&self, u: NonZeroUsize, v: NonZeroUsize) -> NonZeroUsize {
        let mut u = u.get();
        let mut v = v.get();
        if self.distances[u] < self.distances[v] {
            std::mem::swap(&mut u, &mut v);
        }

        let max_step = self.parents_by_step.len();
        for k in 0..max_step {
            if (self.distances[u] - self.distances[v]) >> k & 1 == 1 {
                u = self.parents_by_step[k][u].unwrap().get();
            }
        }
        if u == v {
            return NonZeroUsize::new(u).unwrap();
        }
        for k in (0..max_step).rev() {
            if self.parents_by_step[k][u] != self.parents_by_step[k][v] {
                u = self.parents_by_step[k][u].unwrap().get();
                v = self.parents_by_step[k][v].unwrap().get();
            }
        }
        self.parents_by_step[0][u].unwrap()
    }

    /// Calculates the distance between the vertex `u` and `v`.
    pub fn dist(&self, u: NonZeroUsize, v: NonZeroUsize) -> usize {
        let root_dist = self.distances[self.query(u, v).get()];
        let u = u.get();
        let v = v.get();
        self.distances[u] + self.distances[v] - 2 * root_dist
    }
}