そらで書ける <Θ(N), Θ(log N)> Level Ancestor

次の問題を考えます.

問題 (Level Ancestor Problem)

頂点 $1$ を根とする $N$ 頂点の木 $T$ が与えられます.

頂点 $v$ と非負整数 $k$ に対して,$\mathrm{LA}(v, k)$ を頂点 $v$ から親方向への辺をちょうど $k$ 回辿って到達する頂点と定義します.

頂点と非負整数の組が $Q$ 個与えられます.$i$ 番目の組を $(v _ i, k _ i)$ として,各 $i$ に対して $\mathrm{LA}(v _ i, k _ i)$ を求めて出力してください.ただし,頂点 $v _ i$ から親方向への辺を $k _ i$ 回辿れない場合は,代わりに $-1$ を出力してください.

本記事の意義 (?)

本節は読み飛ばして貰っても構いません.

Level Ancestor Problem の解法として,Heavy Light Decomposition (HLD) を用いたものや,ダブリングによるものがよく用いられています (要出典).それぞれの解法について詳しくは述べませんが,計算量は次のようになっています.

前計算 クエリ 空間
HLD $\Theta(N)$ $\Theta(\log N)$ $\Theta(N)$
ダブリング $\Theta(N\log N)$ $\Theta(\log N)$ $\Theta(N \log N)$
本記事で紹介するアルゴリズム $\Theta(N)$ $\Theta(\log N)$ $\Theta(N)$

上表にも示しましたが,本記事では Level Ancestor Problem を前計算 $\Theta(N)$,クエリ $\Theta(\log N)$ で解くアルゴリズムを紹介します.表を見ると HLD を用いた解法で十分そうに見えます.実際 HLD で十分なのですが,Level Ancestor Problem を解くためだけに HLD を実装するのは,やや大変です.実装の簡単さの点ではダブリングによる解法は優れていますが,前計算に掛かる時間計算量や空間計算量のオーダーが HLD のそれと比較して悪化してしまいます.本記事では,HLD による解法と同じ計算量オーダーを達成しつつ,実装も簡単なアルゴリズムを紹介します.

なお,クエリ先読みが可能な場合 (Offline Level Ancestor Problem) は,以下の記事で紹介されている手法で解くのが楽だと思います (計算量の面でも,優れています).

noshi91.hatenablog.com

本題

書き切ってから気付いたんですが,先に「おわりに」の節を読んだ方が何をしようとしているかわかりやすいかもしれません.


クエリ v k の処理を考えましょう.頂点 $v$ の深さを $\mathrm{de p}(v)$ とすると,$\mathrm{LA}(v, k)$ が存在するための必要十分条件は $\mathrm{de p}(v) \geq k$ を満たすことです.そこで,以降は $\mathrm{de p}(v) \geq k$ を仮定します.

深さがちょうど $d$ であるような頂点の集合を $S _ d$ とすると,$\mathrm{LA}(v, k) \in S _ {\mathrm{de p}(v) - k}$ が成り立ちます.また,$S _ {\mathrm{de p}(v) - k}$ に含まれる頂点のうち,$v$ の先祖 (ここでは,$v$ 自身も含みます) であるようなものはちょうど $1$ つしか存在しないので,$S _ {\mathrm{de p}(v) - k}$ に含まれる $v$ の先祖を高速に見つけられればよいです.しかし,$S _ {\mathrm{de p}(v) - k}$ のサイズは最大で $\Theta(N)$ となるため,線形探索では最悪の場合は $\Omega(N)$ time 掛かってしまいます.そこで,上手く探索出来るように工夫をします.

木 $T$ に対して根 $1$ を始点とする深さ優先探索を行い,その訪問順 (pre-order) に頂点を並べ替えてできる頂点の列を $v _ 1, \ldots, v _ N$ とします.そして,$v _ i$ の頂点番号が $i$ となるように頂点番号を振り直します.古い頂点番号から新しい頂点番号への変換やその逆変換は各順列を覚えておくことで定数時間で行えます.そこで,以降は新しく振り直した頂点番号を用いて考えます.

頂点番号の振り直しの例

さて,このような頂点番号の振り方をすることで,頂点 $u$ を根とする部分木のサイズを $\mathrm{sub}(u)$ とすれば,頂点 $u$ を根とする部分木は半開区間 $[u,u+\mathrm{sub}(u))$ を成します.この事実から,$\mathrm{LA}(v,k)$ を根とする部分木が成す区間を考えることで, $\mathrm{LA}(v,k) \leq v$ を得ます.

ここで,$S _ {\mathrm{de p}(v) - k}$ の要素 $x$ であって $x \leq v$ を満たす頂点のうち最大のものを $p$ とすれば,実は $\mathrm{LA}(v,k) = p$ が成り立ちます $(\star)$.従って,$S _ {\mathrm{de p}(v) - k}$ の要素を昇順に並べた列を予め計算しておくことで,二分探索により $\mathrm{LA}(v,k)$ を $\Theta(\log N)$ time で計算することが出来ます.

$(\star)$ の証明

まず,$\mathrm{LA}(v,k) \leq v$ より,$p$ は必ず存在します.

$p$ 以外に $x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ が存在しない場合は,明らかに $\mathrm{LA}(v,k) = p$ が成り立ちます.

$p$ 以外に $x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ が存在する場合を考えます.$x \leq v$ を満たす頂点 $x \in S _ {\mathrm{de p}(v) - k} \backslash \{p\}$ を任意に取り,$x = \mathrm{LA}(v,k)$ を仮定します.このとき,$x$ は $v$ の先祖なので,$v \in [x, x + \mathrm{sub}(x))$ が成り立ちます.ところで $p$ は $x$ を根とする部分木には含まれず,また $p$ の最大性から $x\lt p$ が成り立つため,$p \geq x + \mathrm{sub}(x)$ を満たす必要があります.従って,$v \in [x, x + \mathrm{sub}(x))$ より $x \leq v \lt x + \mathrm{sub}(x) \leq p$,即ち $v\lt p$ を得,これは $p \leq v$ であることに矛盾します.

以上より,$p = \mathrm{LA}(v, k)$ が成り立ちます.(証明終)

各 $d$ に対して $S _ d$ を昇順に並べた列 $A _ d$ を計算する必要がありますが,これは深さ優先探索において頂点 $u$ を訪問したタイミングで $A _ {\mathrm{de p}(u)}$ の末尾に $u$ を追加することで,前計算に掛かる計算量を $\Theta(N)$ に保ったまま構築することが出来ます.

以上より,Level Ancestor Problem を前計算 $\Theta(N)$ time,クエリ $\Theta(\log N)$ time,空間 $\Theta(N)$ で解くことが出来ました.

実装

折りたたみ

#include <algorithm>
#include <vector>

struct LevelAncestor {
    LevelAncestor() = default;
    // 隣接リスト g と根 root 
    LevelAncestor(const std::vector<std::vector<int>> &g, const int root = 0) : _n(g.size()), _visit_time(_n), _depth(_n), _bucket(_n) {
        build(g, root);
    }
    
    // u から親を k 回辿って到達する頂点を返す.0<=k<=dep[u] を満たさない場合は,そのような頂点は存在しないので -1 を返す.
    int query(const int u, const int k) const {
        const int d = _depth[u] - k;
        if (d < 0 or d > _depth[u]) return -1;
        // 頂点 i と頂点 j を訪問時間で比較する関数
        auto comp = [this](int i, int j) {
            return _visit_time[i] < _visit_time[j];
        };
        // u 以下のうち最大 = u より真に大きい頂点のうち最小のものの 1 つ前
        // なので,upper_bound をしてから prev で 1 つ戻せばよい
        return *std::prev(std::upper_bound(_bucket[d].begin(), _bucket[d].end(), u, comp));
    }
    int operator()(const int u, const int k) const {
        return query(u, k);
    }

private:
    // 頂点数
    int _n;
    // 深さ優先探索における訪問時刻
    std::vector<int> _visit_time;
    // 頂点の深さ
    std::vector<int> _depth;
    // 記事中の S.各々訪問時刻の昇順にソートされている
    std::vector<std::vector<int>> _bucket;

    void build(const std::vector<std::vector<int>> &g, const int root) {
        // 現在時刻
        int time = 0;
        // u : 現在見ている頂点
        // p : u の親
        auto dfs = [&](auto dfs, int u, int p) -> void {
            // 頂点 u の訪問時刻を記録して,時刻を 1 つ進める
            _visit_time[u] = time++;
            // S _ {dep(u)} に u を追加
            _bucket[_depth[u]].push_back(u);
            for (int v : g[u]) {
                if (v == p) continue;
                _depth[v] = _depth[u] + 1;
                dfs(dfs, v, u);
            }
        };
        dfs(dfs, root, -1);
    }
};

(2022/06/08 追記)

上記の実装は std::upper_bound に渡した比較器内で配列アクセスを行っているため,遅いです.訪問時刻をバケットに詰めて配列アクセスの回数を減らすことで定数倍高速化を図ることが出来ます.

以下に高速な実装例を示しておきます.(そらで書けると謳っておきながら,そこそこ面倒くさい感じになってしまいました)

suisen-cp.github.io

(追記終)

検証

Library Checker

Level Ancestor を verify できる問題が思い浮かばなかったので,LA で二分探索することで LCA を計算するコードを書いて,LCA を verify する問題に投げています.

おわりに

正当性を示すために長々と書いてしまいましたが,アルゴリズムの中身はかなり直感的で実装もしやすいと思います.やっていることは,前計算では DFS の訪問順に深さ毎に用意したバケツに頂点を詰め込んで,クエリでは対応するバケツを見て訪問時刻が自身の訪問時刻を超えないもののうち最も遅く訪れた頂点を二分探索で求めるだけなので.