트리 (2) - 희소 배열, 최소 공통 조상(LCA)

 

 

개미(BOJ 14942)를 보자. 문제 자체는 정말 단순한데, 각 정점별로 dfs하는 간단한 해법을 떠올릴 수 있다. 그러나 트리가 일자로 깊은 구조라면, $O(N^2)$의 시간이 소요되어 시간 초과를 받게 될 것이다. 따라서 depth가 깊은 경우 빠르게 조상으로 거슬러 올라가는 방법이 필요하다.

 

 

희소 배열

희소 배열(sparse table) 자료 구조가 이 문제를 해결해 준다. 아이디어는 매우 간단한데, 각 노드의 부모를 저장하는 parent 배열을 2차원으로 확장시켜 다음과 같이 정의하는 것이다.

 

$parent[i][k]$ = ($i$번 노드의 $2^k$번째 조상)

 

부모는 1번째 조상이므로 k=0에 해당한다. 2번째 조상을 구할 때는 부모의 부모를 찾으면 되며, 4번째 조상은 2번째 조상의 2번째 조상일 것이다. 따라서 점화식을 쉽게 구성할 수 있고, k를 증가시켜 가며 희소 배열을 채울 수 있다.

for k in range(1, 16):
    for i in range(1, N+1):
        parent[i][k] = parent[parent[i][k-1]][k-1]

 

 

희소 배열을 완성한 후에는, 거꾸로 큰 k부터 시작하여 $2^k$번째 조상으로 점프하는 것이 가능하면 점프하면 된다. 로그 시간이 걸릴 것이므로 아까 다룬 문제의 경우 $O(N \log N)$만에 해결이 가능하다.

 

사실 희소 배열은 원본 배열이 업데이트되지 않는 조건 하에 트리 외 다양한 상황에도 접목이 가능하다. 대표적으로 배열 내 구간의 최댓값이나 최솟값을 구하는 경우, $[i, i+2^k]$ 구간을 희소 배열로 전처리하여 쿼리당 $O(1)$ 처리가 가능하다. 희소 배열의 또다른 조건은 합성함수 조건 $f^{(a+b)} = f^a \circ f^b $이다. 이 조건이 성립해야 임의의 $D$를 비트 단위로 쪼개 이동하여 총 $D$만큼 이동하는 것이 가능하기 때문이다. 

 

 

최소 공통 조상

최소 공통 조상(least common ancestor)은 이름 그대로 트리 내 두 정점의 최소 공통 조상을 찾는 방법이다. 기본적인 알고리즘은 다음과 같다.

 

1. 두 정점의 depth를 구한다. 이는 루트 노드로부터 dfs/bfs하여 depth배열을 전처리해 둠으로써 가능하다.
2. 두 정점 중 depth가 더 깊은 정점을 depth가 같아질 때까지 올린다.
3. 각 정점을 하나씩 올려가며 최소 공통 조상이 나올 때까지 반복한다.

 

이를 naive하게 구현하면 2와 3에서 $O(N)$이 소요된다. 희소 배열을 이용하여 아래와 같이 구현하면, $O(\log N)$ 풀이가 가능하다.

a, b = map(int, input().split())
    if depth[a] > depth[b]: a, b = b, a
    a_old, b_old = a, b
    diff = depth[b] - depth[a]
    for i in range(15, -1, -1):
        bit = (diff >> i) & 1
        if bit: b = parent[b][i]

    if a == b: lca = a
    else:
        k = 15
        while k >= 0:
            pa = parent[a][k]
            pb = parent[b][k]
            if pa == pb: k -= 1
            else:
                a = pa
                b = pb
        lca = pa

 

 

최소 공통 조상은 트리 내 두 정점 사이의 거리를 구하는 데 사용될 수 있다. 간선의 중복 없이 트리 내 두 정점을 잇는 경로는 lca를 거치는 경로로 유일하므로, 다음과 같이 구하면 된다.

 

$$dist(a, b) = depth[a] + depth[b] - 2depth[lca]$$