트리 (5) - Heavy-light decomposition

센트로이드 분할에 이어, 트리의 경로 쿼리를 로그 시간에 처리할 수 있도록 해 주는 또 다른 알고리즘인 Heavy-light decompositon(HLD)에 대해 알아보자.

 

 

HLD의 원리

hld의 기본적인 아이디어는 Eular tour technique(ETT)처럼 트리에 dfs order를 매겨 트리를 일자로 편 뒤, 쿼리를 세그먼트 트리 상에서 처리하자는 것이다. 다만 서브트리 전체에 관한 쿼리가 들어오는 경우 ett로 쉽게 처리가 가능하지만, 특정 두 정점 사이의 경로에 관한 쿼리의 경우 어떻게 dfs order를 매길 것인지, 어떻게 세그먼트 트리에서 쿼리를 처리할 것인지 쉽게 떠올리기 어렵다. hld는 이러한 의문에 대한 효율적인 정답을 제시한다.

 

먼저 dfs 순회 방법에 관하여, hld는 서브트리의 size가 제일 큰 정점으로 이어지는 간선을 heavy edge, 그렇지 않은 나머지 간선들을 light edge로 분류한다. 이 중 heavy edge에 대해서 먼저 dfs하고 나머지는 그 이후에 돌면서 dfs order를 매겨준다. 정점을 내려가면서 이 과정을 반복하면, 끝에 도달할 때까지 계속해서 heavy edge로 먼저 이동할 것이므로 heavy edge들로 구성된 heavy-chain이 dfs 순회 결과 연속적으로 나타날 것이다.

 

https://www.geeksforgeeks.org/dsa/introduction-to-heavy-light-decomposition/

 

가령 위의 그래프에서, dfs 순회 결과는 [1,2,6,8,10,11,5,3,7,9,4]와 같이 나타날 수 있다. 따라서 찾는 경로가 heavy-chain에 속하는 경우라면, 세그먼트 트리 상에서 구간 쿼리를 날려서 쉽게 처리할 수 있을 것이다.

 

한편 찾는 경로가 light edge를 포함하는 경우, 다음과 같이 진행하면 된다.

1. dfs order 순서 $id[u]$와, 각 정점들이 속한 heavy-chain의 가장 위쪽 원소를 저장하는 배열 $top[u]$를 미리 구해 놓는다.
2. $top[u], top[v]$ 중 깊이가 더 깊은 정점을 대상으로(편의상 $u$라고 가정), $(id[top[u]], id[u])$까지 구간 쿼리를 날린다.
3. $u := parent[top[u]]$로 갱신한다.
4. $top[u] = top[v]$가 될 때까지 2-3을 반복한다.
5. 마지막으로, 같은 heavy-chain에 속하게 된 $u, v$를 대상으로 구간 쿼리를 날려 준다.

light 간선을 건너는 순간 서브트리의 크기는 절반 이하로 떨어짐이 보장되므로, 쿼리 처리 시 light 간선을 건너는 횟수는 최대 $O(\log N)$이다. 따라서 세그먼트 트리 연산까지 총 $O(\log^2 N)$만에 쿼리를 처리할 수 있다는 것이 hld의 핵심이다.

 

 

HLD의 구현

첫 번째 dfs에서 depth, parent, size를 구해주고, size에 따라 정렬해 heavy 간선이 맨 처음 위치하도록 옮겨 준 뒤, 두 번째 dfs에서 heavy 간선 먼저 돌면서 id와 top을 기록해 주면 된다.

size = [1]*(N+1)
parent = [0]*(N+1)
depth = [0]*(N+1)
def dfs_size(u, p, d):
    parent[u] = p
    depth[u] = d
    heavy_size = 0
    heavy_idx = 0
    for i, v in enumerate(adj[u]):
        if v == p: continue
        dfs_size(v, u, d + 1)
        size[u] += size[v]
        if heavy_size < size[v]:
            heavy_size = size[v]
            heavy_idx = i
    adj[u][0], adj[u][heavy_idx] = adj[u][heavy_idx], adj[u][0]

id = [0]*(N+1)
top = [0]*(N+1)
count = 0
def dfs_hld(u, p, t):
    global count
    count += 1
    id[u] = count
    top[u] = t
    
    # heavy
    if adj[u][0] != p: dfs_hld(adj[u][0], u, t)
    
    # light
    for v in adj[u][1:]:
    	if v != p: dfs_hld(v, u, v)

 

 

물론 파이썬의 재귀 dfs는 매우 느리기 때문에 스택을 활용할 수도 있다. 이 경우 heavy를 먼저 순회하도록 하려면 heavy를 나중에 스택에 투입해야 한다.

size = [1]*(N+1)
dep = [0]*(N+1)
parent = [0]*(N+1)
stack = [1]
visit_order = []
while stack:
    u = stack.pop()
    visit_order.append(u)
    for v, i in adj[u]:
        if v == parent[u]: continue
        parent[v] = u
        dep[v] = dep[u] + 1
        stack.append(v)

for u in reversed(visit_order): 
    heavy_size = 0
    heavy_idx = 0
    for i, v in enumerate(adj[u]):
        if v == parent[u]: continue
        size[u] += size[v]
        if heavy_size < size[v]:
            heavy_size = size[v]
            heavy_idx = i
    adj[u][0], adj[u][heavy_idx] = adj[u][heavy_idx], adj[u][0]


id = [0]*(N+1)
top = [0]*(N+1)
count = 0
stack = [(1, 0, 1)]
while stack:
    u, p, t = stack.pop()
    count += 1
    id[u] = count
    top[u] = t

    for i in range(len(adj[u])-1, 0, -1):
        v = adj[u][i]
        if v != p: stack.append((v, u, v))
    
    v = adj[u][0]
    if v != p: stack.append((v, u, t))

 

 

다음으로 쿼리 처리를 해준다. 트리와 쿼리 1(BOJ 13510)처럼 경로 중 가중치가 최대인 간선을 찾는 쿼리가 들어온다고 해 보자. 세그먼트 트리에 인덱스 $id[u]$가 $(parent[u], u)$를 잇는 간선의 가중치를 가리키도록 저장해 두고, 경로를 heavy-chain 단위로 잘 쪼개서 구간 쿼리를 날려 최댓값을 갱신해 주면 된다. 마지막에 $u, v$가 같은 chain 안에 있을 경우, ($dep[u] < dep[v]$라 가정하면) u의 부모 간선은 포함하면 안 되므로 $(id[u]+1, id[v])$로 구간 쿼리를 날려야 함에 주의하자.

res = 0
while top[u] != top[v]:
    if dep[top[u]] < dep[top[v]]: u,v = v,u
    res = max(res, tree.query(1, 1, N, id[top[u]], id[u]))
    u = parent[top[u]]
if dep[u] > dep[v]: u,v = v,u
res = max(res, tree.query(1, 1, N, id[u]+1, id[v]))
print(res)

 

 

센트로이드와 비교해 보면, 센트로이드는 정점을 효율적으로 재배치하여 다양한 경로 쿼리에 대응 가능한 반면에(ex. 한 정점부터 모든 정점까지의 경로, 임의의 두 정점 간 경로), 기본적으로 ett와 비슷한 원리인 hld는 특정 두 정점 간 경로의 처리에 특화되어 있다. hld는 세그먼트 트리로 동작하므로 lazy propagation 등이 가능하다는 장점을 갖고, 코드는 비교적 길지만 알고리즘이 직관적이어서 이해 및 응용하기에 보다 수월할 수 있다.