Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

LCAを実装した

LCA (Lowest Common Ancestor)は、木構造の中でもっとも低い共通の祖先を調べるデータ構造です。

これはナイーブに実装すると、ノード数分だけ時間がかかってしまいますが、Sparse Tableと同様のアイデアを用いることによってO(logN)まで落とすことが出来ます。

イデアは、

  • Sparse Tableの木版のようなかんじで、自分の2k上の親をテーブルに保存します。各ノードの深さもテーブルに保存しておきます
  • 2つのノードvが与えられた時、共通の祖先までの距離を2進数で表せます。従って、作ったテーブルを使って何度かO(logN)回ジャンプすれば届くことはわかります
  • u,vのLCAを求める場合は、まずu,vの深さが同じになるまで深い方のノードを上にジャンプします。これは、O(logN)回で済みます
  • あとは、大きなジャンプから試して、お互いにLCAの一歩手前までたどり着きます(ここがうまい。このアイデアによって「ジャンプした先が同じにならない限りはジャンプする」という実装が可能になる)。そして、その一つ上が求めるノードになります

asssertだらけなのでわかると思いますが、これもやはりハマりまくったため、コーヒーショップでFワードを連発しながら実装しました。まぁ、assert自体は可読性に役立ってるので良いと思いますが。

struct LCA <'a> {
    root: usize,
    tree: &'a [Vec<usize>],
    parent: Vec<Vec<Option<usize>>>,
    depth: Vec<u32>,
}

impl <'a> LCA<'a> {
    fn new(root: usize, tree: &'a [Vec<usize>]) -> Self {
        let n = tree.len();
        let mut log_n = (n as f64).log2().ceil() as usize;
        if log_n == 0 {
            log_n = 1;
        }
        assert!(log_n > 0);
        Self {
            root,
            tree,
            parent: vec![vec![None; n]; log_n],
            depth: vec![0; n],
        }
    }
    // store direct parent and depth
    fn dfs(&mut self, u: usize, parent: Option<usize>, depth: u32) {
        self.parent[0][u] = parent;
        self.depth[u] = depth;
        for i in 0 .. self.tree[u].len() {
            let v = self.tree[u][i];
            if Some(v) != parent {
                self.dfs(v, Some(u), depth+1);
            }
        }
    }
    fn build(&mut self) {
        let root = self.root;
        self.dfs(root, None, 0);

        let mut k = 0;
        while k+1 < self.parent.len() {
            for u in 0 .. self.tree.len() {
                if self.parent[k][u].is_some() {
                    self.parent[k+1][u] = self.parent[k][self.parent[k][u].unwrap()]
                } 
            }
            k += 1;
        }
    }
    fn lca(&self, u: usize, v: usize) -> usize {
        let (mut v0, mut v1) = if self.depth[u] <= self.depth[v] {
            (u, v)
        } else {
            (v, u)
        };
        assert!(self.depth[v1] >= self.depth[v0]);

        // move v1 up until depth of v0 and v1 gets equal.
        for k in 0 .. self.parent.len() {
            if (((self.depth[v1] - self.depth[v0]) >> k) & 1) > 0 {
                assert!(self.parent[k][v1].is_some());
                v1 = self.parent[k][v1].unwrap();
            }
        }
        assert!(self.depth[v1] >= self.depth[v0]);
        assert!(self.depth[v1] == self.depth[v0]);
        if (v0 == v1) {
            return v0;
        }
        for k in (0..self.parent.len()).rev() {
            // LCA's parent is LCA
            if self.parent[k][v0] != self.parent[k][v1] {
                assert!(self.parent[k][v0].is_some());
                assert!(self.parent[k][v1].is_some());
                v0 = self.parent[k][v0].unwrap();
                v1 = self.parent[k][v1].unwrap();
            }
        }
        return self.parent[0][v0].unwrap();
    }
}

他には、木構造をDFSでたどって、必ずu,vという順で現れるように木構造を配列に落とし込んでしまえば、Sparse Tableを使ってRMQすれば終わるよねというやり方もあります(LCAはuとvに挟まれたもっとも深さの小さいノードなので)。