GRL_5_C - Lowest Common Ancestor / PythonでAizu Online Judgeを解く

Posted on 2019/07/15

LCA: 最小共通祖先

根付き木の2頂点について最も根から遠い共通の祖先となる頂点(LCA)を求めてください。木は \(n\) 個の頂点からなり各頂点には根を \(0\) として \(0\) から \(n-1\) までの ID がそれぞれ割り振られるものとします。

入力

%python
def input(f):
    n, = map(int, f.readline().split())
    k = [None for _ in range(n)]
    c = [[] for _ in range(n)]
    for i in range(n):
        tmp = map(int, f.readline().split())
        k[i] = tmp.pop(0)
        c[i] = tmp[:]
    q, = map(int, f.readline().split())
    u = [None for _ in range(q)]
    v = [None for _ in range(q)]
    for i in range(q):
        u[i], v[i] = map(int, f.readline().split())
    
    return n, k, c, q, u, v

制約

  • \(1 \leq n \leq 100,000\)
  • \(1 \leq q \leq 100,000\)

サンプル入力

%sh
cat <<EOF > /tmp/input1.txt
8
3 1 2 3
2 4 5
0
0
0
2 6 7
0
0
4
4 6
4 7
4 3
5 2
EOF

%python
with open('/tmp/input1.txt') as f:
    print(input(f))
(8, [3, 2, 0, 0, 0, 2, 0, 0], [[1, 2, 3], [4, 5], [], [], [], [6, 7], [], []], 4, [4, 4, 4, 5], [6, 7, 3, 2])

考え方

🐕 とりあえず LCA を求めたい

やりたいことは結構単純で指定された二つの頂点から根に向かって移動するときに「一番最初に合流可能な頂点」がどこになるかという話です。

アプローチとしても二つの頂点を起点にまずは深さを揃えてあげて共通の頂点に到達するまで同時に移動すればよく計算量としても高々頂点数程度になるかと思います。

ただ、ここでは \(10^5\) 個ものクエリを捌く必要があり \(\mathcal O (n)\) だと制限時間には間に合いません。各クエリの処理を \(\mathcal O (log\ n)\) 程度まで抑えるために Doubling と呼ばれる方法が利用できます。

🐕 Doubling とは

  • \(pa(i, j)\) を頂点 \(j\) から根頂点に向かって \(2^i\) 番目にある頂点とします
    • ここで \(pa(i+1, j) = pa(i, pa(i, j))\) となります
      • 親要素の親要素というイメージ
  • また、\(i = 0\) のとき \(pa(0, j)\) は頂点 \(j\) の親要素であるものとします
  • \(k = log_2(n)\) を \(i\) の上限として \(0\) から \(k\) までの \(i\) について \(pa(i+1, j) = pa(i, pa(i, j))\) を計算します
    • \(i+1\) を求めるとき \(pa(i, j)\) は計算済みです

🐕 クエリの捌き方

  • \(pa\) とは別に頂点の深さを表す \(depth\) も前計算しておきます(\(depth(i)\) を頂点 \(i\) の深さとします)
  • LCA を求める頂点 \(u, v\) について \(depth[u] = depth[v]\) を満たすまで以下の処理を繰り返します
    • \(d = depth[u] - depth[v]\) とします
      • ※ \(depth[v] < depth[u]\) を満たさない場合は \(u\) と \(v\) を入れ替えます
    • \(d\) のバイナリ表現について \(x\) ビット目が 1 になっていたら \(v = pa[x][v]\) として v を更新します
  • 高さが揃ったら共通の頂点にたどり着くまで一段ずつ \(u = pa[k][u]\) と \(v = pa[k][v]\) を辿っていきます

上記のような手順で LCA を効率的に求めることが可能です。前計算で \(pa\) の構築が \(\mathcal O(n\ log\ n)\) 程度、クエリの処理が \(\mathcal O(log\ n)\) 程度に収まりました。

実装

%python
import math
from collections import deque

def solve(n, k, c, q, u, v):
    graph = [[] for _ in range(n)]
    for i in range(n):
        for j in range(k[i]):
            graph[i].append(c[i][j])
    
    logn = int(math.log(n, 2))+1
    pa = [[None for _ in range(n)] for _ in range(logn+1)]
    depth = [None for _ in range(n)]
    
    def build():
        que = deque()
        que.append(0)
        depth[0] = 0
        while que:
            cur = que.popleft()
            for child in c[cur]:
                depth[child] = depth[cur] + 1
                que.append(child)
            
        for i in range(n):
            for child in c[i]:
                pa[0][child] = i
        for i in range(logn):
            for j in range(n):
                if pa[i][j] != None:
                    pa[i+1][j] = pa[i][pa[i][j]]
    
    def query(u, v):
        def move(src, dst):
            return src, dst
        
        while depth[u] != depth[v]:
            if depth[v] < depth[u]:
                u, v = v, u
            d = depth[v] - depth[u]
            for k in range(logn):
                if d & 1:
                    move(v, pa[k][v])
                    v = pa[k][v]
                d = d >> 1
        for nk in range(logn):
            k = logn - nk - 1
            if pa[k][u] != pa[k][v]:
                move(u, pa[k][u])
                u = pa[k][u]
                move(v, pa[k][v])
                v = pa[k][v]
        if u == v:
            res = u
        else:
            move(u, pa[0][u])
            res = pa[0][u]
        
        return res
    
    depth[0] = 0
    build()

    for i in range(q):
        yield query(u[i], v[i])

出力サンプル

%sh
for x in 3 5 6 11; do
    curl -s https://judgedat.u-aizu.ac.jp/testcases/GRL_5_C/$x/in > /tmp/input$x.txt
done

%python
with open('/tmp/input5.txt') as f:
    for res in solve(*input(f)):
        print(res)
4
4
4
3
1
1
1
1
2
3

Wrong Answers 😇

  • Doubling するときの上限設定ミス
    • log2 じゃなくて自然対数になっていました😇
  • 高さが揃ったあとの動かし方
    • 根に向かって同時移動させるときは pa の i を上限から動かしましょう

アニメーション

猫がジャンプしている様子が見受けられます :-)