HL分解 (Heavy-Light Decomposition) をRustで実装した

はじめに

木まわりの問題で毎回 LCA やパス上のクエリを書き直してバグらせまくっていたので、
Rustで Heavy-Light Decompositionをちゃんとライブラリ化しました。
せっかくなので、自分用メモも兼ねて記事にしました。

できること

Heavy-Light Decompositionは、木のパスをいくつかの連続した区間に分解することで計算量をサボれるアルゴリズムです。
この実装では、次のようなことがO(log N) でできます。

lca(u, v)

頂点 u, v の LCA を返す

distance(u, v)

u, v 間の距離

go(from, to, times)

from → to のパス上で、from から times 個進んだ頂点を返す（存在しなければ None）

foreach(u, v, f)

u-v パスを HLD の順番 (vid) 上の区間 [l, r] に分解しながら、各区間に対して f(l, r) を呼ぶ
*当然ですがここにセグ木や Fenwick Tree をかませる場合はO(log^2N)になりますね。

アルゴリズム概要（なんで O(log N) になるのか）

HL分解では、まず O(N) で DFS して各頂点の部分木サイズを求め、いちばん大きい子を heavy[v] として記録します。
v → heavy[v] の辺を heavy edge, その他の子への辺を light edge と呼びます。

light edge が高々 O(log N) 本しか出てこない

ある頂点 v の部分木サイズを sz(v) とします。
heavy[v] は子の中で一番大きいので、light child c に対しては

• sz(c) ≤ sz(v) / 2

が成り立ちます（最大の子より小さいので、合計しても sz(v) を超えない）。

つまり light edge を 1 本進むたびに、部分木サイズは「少なくとも半分以下」になります。

root → 任意の頂点 x のパスを考えると、
sz(root) = N から始まって、light edge を通るたびに

• N → N/2 → N/4 → …

と 1/2 ずつ減っていくので、light edge は高々 log₂ N 本しか現れません。

直感的には平衡二分木で根から葉まで行くときに高さが O(log N) になるイメージで理解しています

2 回目の DFS でチェインに番号を振る

次に、dfs_hld でもう一度 DFS します。このとき

• heavy edge を優先して辿るように再帰し、
• heavy パスごとに vid が連続した区間 [l, r] になるように番号を振ります。
• 同時に、そのパスの先頭を head[v] に記録します。

これで

• 「同じチェイン上にある頂点たち」は vid 上で連続した区間になり、
• 「チェインの先頭」は head[v] として 1 つ上に飛べる

という状態が作られます。

各操作の計算量

• lca(u, v)
head[u] != head[v] の間、深い方のチェインの先頭の親にジャンプしています。
  このジャンプは「light edge を 1 本だけ遡る」のと同じなので、
  root→u, root→v に含まれる light edge の本数 O(log N) 回で必ず終わります。

• distance(u, v)
lca(u, v) を 1 回呼んで
  depth[u] + depth[v] - 2 * depth[lca] を返すだけなので O(log N)。

• ancestor(v, k) / go(from, to, times)
  どちらも「同じチェイン内で上れるだけ上って、足りなければチェインの先頭の親へジャンプ」を繰り返すので、
  やはり light edge の本数ぶんしかループが回らず O(log N)。

• foreach(u, v, f)
  LCA と同じように、チェインごとに区間 [l, r] に分解しながら
  深い方のチェインの先頭の親へジャンプしていくので、
  呼ばれる区間の個数は高々 O(log N) 個です。
  ここにセグ木（クエリ O(log N)）を乗せると、パスクエリ全体は O(log^2 N) になります。

アルゴリズムの詳解はこちらの記事がすごく丁寧でわかりやすいです。

実装

実装です。

pub struct HLDecomposition {
    n: usize,
    root: usize,
    g: Vec<Vec<usize>>,
    parent: Vec<usize>,
    depth: Vec<usize>,
    heavy: Vec<usize>,
    // 頂点番号の再割り当ての際に使う変数
    t: usize,
    // 頂点番号を再割り当てした際のindex
    vid: Vec<usize>,
    // vidの逆
    inv: Vec<usize>,
    // dfsに入って来たときのtの値
    t_in: Vec<usize>,
    // dfs戻るときのtの値
    t_out: Vec<usize>,
    // heavy edgeの親
    head: Vec<usize>,
    is_built: bool,
}

impl HLDecomposition {
    pub fn new(n: usize) -> Self {
        Self {
            n,
            root: usize::MAX,
            g: vec![vec![]; n],
            parent: vec![usize::MAX; n],
            depth: vec![usize::MAX; n],
            heavy: vec![usize::MAX; n],
            t: 0,
            vid: vec![usize::MAX; n],
            inv: vec![usize::MAX; n],
            t_in: vec![usize::MAX; n],
            t_out: vec![usize::MAX; n],
            head: (0..n).collect::<Vec<usize>>(),
            is_built: false,
        }
    }
    pub fn add(&mut self, u: usize, v: usize) {
        self.g[u].push(v);
        self.g[v].push(u);
    }

    pub fn build(&mut self, root: usize) {
        assert!(self.n > 0);
        self.root = root;
        self.depth[root] = 0;
        self.dfs(root, usize::MAX);
        self.t = 0;
        self.head[root] = root;
        self.dfs_hld(root);
        self.is_built = true;
    }

    pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
        debug_assert!(self.is_built);
        while self.head[u] != self.head[v] {
            if self.depth[self.head[u]] > self.depth[self.head[v]] {
                u = self.parent[self.head[u]];
            } else {
                v = self.parent[self.head[v]];
            }
        }
        if self.depth[u] < self.depth[v] {
            u
        } else {
            v
        }
    }

    pub fn distance(&self, u: usize, v: usize) -> usize {
        debug_assert!(self.is_built);
        let l = self.lca(u, v);
        self.depth[u] + self.depth[v] - 2 * self.depth[l]
    }

    /// from → to のパス上で from から times 個進んだ頂点
    pub fn go(&self, from: usize, to: usize, times: usize) -> Option<usize> {
        debug_assert!(self.is_built);
        let d = self.distance(from, to);
        if times > d {
            return None;
        }
        let lc = self.lca(from, to);
        if lc == to {
            return self.ancestor(from, times);
        }
        if lc == from {
            return self.child(from, to, times);
        }
        let d_to_lca = self.distance(from, lc);
        if times < d_to_lca {
            self.ancestor(from, times)
        } else if times == d_to_lca {
            Some(lc)
        } else {
            let rest = times - d_to_lca;
            self.child(lc, to, rest)
        }
    }

    /// u-v パスを (vid 上の) 区間 [l, r] に分解して f(l, r) を呼ぶ
    pub fn foreach<F>(&self, mut u: usize, mut v: usize, mut f: F)
    where
        F: FnMut(usize, usize),
    {
        debug_assert!(self.is_built);
        while self.head[u] != self.head[v] {
            if self.depth[self.head[u]] > self.depth[self.head[v]] {
                std::mem::swap(&mut u, &mut v);
            }
            let hv = self.head[v];
            let l = self.vid[hv];
            let r = self.vid[v];
            f(l, r);
            v = self.parent[hv];
        }

        if self.vid[u] > self.vid[v] {
            std::mem::swap(&mut u, &mut v);
        }
        f(self.vid[u], self.vid[v]);
    }

    // ── getter ──────────────────────────────────────────────────────────
    pub fn g(&self) -> &Vec<Vec<usize>> {
        &self.g
    }
    pub fn parent(&self) -> &Vec<usize> {
        &self.parent
    }
    pub fn depth(&self) -> &Vec<usize> {
        &self.depth
    }
    pub fn heavy(&self) -> &Vec<usize> {
        &self.heavy
    }
    pub fn vid(&self) -> &Vec<usize> {
        &self.vid
    }
    pub fn inv(&self) -> &Vec<usize> {
        &self.inv
    }
    pub fn t_in(&self) -> &Vec<usize> {
        &self.t_in
    }
    pub fn t_out(&self) -> &Vec<usize> {
        &self.t_out
    }
    pub fn head(&self) -> &Vec<usize> {
        &self.head
    }

    // ── internal ────────────────────────────────────────────────────────
    // rootから見た部分木のコスト・深さ最大値・heavy edgeの計算
    fn dfs(&mut self, v: usize, p: usize) -> usize {
        self.parent[v] = p;
        let mut acc = 1;
        let mut m = 0;
        for &v2 in self.g[v].to_owned().iter() {
            if v2 == p {
                continue;
            }
            self.depth[v2] = self.depth[v] + 1;
            let next = self.dfs(v2, v);
            acc += next;
            if next > m {
                m = next;
                self.heavy[v] = v2;
            }
        }
        acc
    }

    fn dfs_hld(&mut self, v: usize) {
        self.vid[v] = self.t;
        self.inv[self.t] = v;
        self.t_in[v] = self.t;
        self.t += 1usize;

        if self.heavy[v] != usize::MAX {
            let next = self.heavy[v];
            self.head[next] = self.head[v];
            self.dfs_hld(next);
        }

        for &v2 in self.g[v].to_owned().iter() {
            if v2 == self.parent[v] || v2 == self.heavy[v] {
                continue;
            }
            self.head[v2] = v2;
            self.dfs_hld(v2);
        }
        self.t_out[v] = self.t;
    }

    // vのk個上の親頂点を返す
    fn ancestor(&self, mut v: usize, mut k: usize) -> Option<usize> {
        while k > 0 {
            let h = self.head[v];
            let d_on_chain = self.depth[v] - self.depth[h];

            if k <= d_on_chain {
                return Some(self.inv[self.vid[v] - k]);
            }

            k -= d_on_chain + 1;
            if self.parent[h] == usize::MAX {
                return None;
            }
            v = self.parent[h];
        }
        Some(v)
    }

    // parent -> childのパスでparentからtimes個だけ進んだ頂点
    fn child(&self, parent: usize, child: usize, times: usize) -> Option<usize> {
        assert!(self.depth[parent] < self.depth[child]);
        let d = self.distance(parent, child);
        if times > d {
            return None;
        }
        let up = d - times;
        self.ancestor(child, up)
    }
}

verify・使い方

yosupoさんのこちらのcheckerで確認できます。

Additiveなセグ木で計算させることで、 前処理をO(N)でできて、クエリ全体ではO(Q * log^2 N)で解けています。

struct M;
impl Monoid for M {
    type S = usize;
    fn identity() -> Self::S {
        0
    }
    fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
        *a + *b
    }
}

fn main() {
    let cin = stdin().lock();
    let mut sc = Scanner::new(cin.bytes().map(|e| e.unwrap() as char));
    let n: usize = sc.get();
    let q: usize = sc.get();
    let a: Vec<usize> = sc.vec(n);

    let mut hld = HLDecomposition::new(n);
    for _ in 0..n - 1 {
        let u: usize = sc.get();
        let v: usize = sc.get();
        hld.add(u, v);
    }
    hld.build(0);

    let mut base = vec![0; n];
    for i in 0..n {
        let &t = hld.vid().get(i).unwrap();
        base[t] = a[i];
    }
    let mut seg: Segtree<M> = Segtree::from(base);
    
    for _ in 0..q {
        let t: usize = sc.get();
        if t == 0 {
            let p: usize = sc.get();
            let x: usize = sc.get();
            let &pre = hld.vid().get(p).unwrap();
            let prenum = seg.get(pre);
            seg.set(pre, x + prenum);
        } else {
            let u: usize = sc.get();
            let v: usize = sc.get();
            let mut ans = 0;
            hld.foreach(u, v, |l, r| {
                ans += seg.prod(l..=r);
            });
            println!("{}", ans);
        }
    }
}

終わりに

バグなど実装改善点などあればコメントください！！！！