백준 11438번: LCA 2 (Python, PyPy3)

접근

이번 문제는 최소공통조상을 구하는데 Sparse Table(희소 테이블)을 이용하는 문제였다.

LCA 알고리즘 개념: 2021.06.27 - [코딩/백준 (Python)] - 백준 3584번: 가장 가까운 공통 조상 (Python)

백준 3584번: 가장 가까운 공통 조상 (Python)

희소 테이블 기본 문제: 2021.06.27 - [코딩/백준 (Python)] - 백준 17435번: 합성함수와 쿼리 (Python)

백준 17435번: 합성함수와 쿼리 (Python)

루트 노드가 1로 주어진 것을 이용하여 bfs 알고리즘을 적용하면 각 노드의 부모 노드가 무엇인지 계산할 수 있다.

이후 희소 테이블 dp를 계산해준다. (이 때, dp[i][j] 는 i 의 2^j 번째 조상을 의미한다.)
2^k = 2^(k - 1) + 2^(k - 1) 이므로, dp[i][j] = dp[ dp[i][j - 1] ][j - 1] 과 같이 나타낼 수 있다.
(i 의 2^(j - 1) 조상의 2^(j - 1) 조상 = i 의 2^j 조상)

이제 주어지는 a 와 b 의 최소공통조상을 구해야 하는데, a 와 b 의 깊이 차이를 보정하여 같은 깊이부터 계산해주기로 한다. 이 때, a 와 b 의 depth 차이가 매우 큰 경우에 노드의 부모를 한칸씩 따라 올라가는 것은 비효율적이다. 이 때도 미리 구한 희소 테이블을 이용하여 계산을 간략히 해줄 수 있다. 특히 비트마스크를 이용해주면 더욱 간단하게 표현할 수 있다.
예를 들어 둘의 깊이 차이가 11 난다고 하면, 이를 이진수로 표현했을 때 1011(2) 이다. 그렇다면 2^3 조상 => 2^1 조상 => 2^0 조상 과정의 세번만 따라가면 11 차이 나는 a 와 b 의 깊이를 맞춰줄 수 있게 된다.

a 와 b 의 깊이를 맞춰준 후, logN 부터, 즉 가장 높은 노드부터 내려오면서 조상들을 비교해주게 된다. 비교를 하다가 dp[a][i] != dp[b][i] 가 되는 지점이 i = k 에서 발생한다면, a 와 b 의 최소공통조상은 a 의 2^k 조상과 2^(k + 1) 조상 사이에 위치하게 된다. 이를 확인하고 난 뒤 a, b = dp[a][i], dp[b][i] 로 갱신시켜 주고, 최초 a 의 2^k 조상과 2^(k + 1) 조상 사이를 계속해서 탐색하게 된다.

코드

import sys
from collections import deque
from math import log2

# tree 정보 입력
n = int(sys.stdin.readline())
tree = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b = map(int, sys.stdin.readline().split())
    tree[a].append(b)
    tree[b].append(a)

# bfs 알고리즘으로 각 노드의 부모노드 및 depth 확인
parent = [0] * (n + 1)
depth = [0] * (n + 1)
check = [0] * (n + 1)
q = deque()
q.append(1)
while q:
    now = q.popleft()
    check[now] = 1
    for i in tree[now]:
        if not check[i]:
            q.append(i)
            parent[i] = now
            depth[i] = depth[now] + 1

# sparse table(dp table) 구하기
logN = int(log2(n) + 1)
dp = [[0] * logN for _ in range(n + 1)]
for i in range(n + 1):
    dp[i][0] = parent[i]
for j in range(1, logN):
    for i in range(1, n + 1):
        dp[i][j] = dp[dp[i][j - 1]][j - 1]

# 비교해야 하는 두 노드의 최소 공통 조상 계산
m = int(sys.stdin.readline())
for _ in range(m):
    a, b = map(int, sys.stdin.readline().split())
    # 깊이 맞추기, 더 깊은 노드 b의 diff 조상을 구하여 깊이를 맞춰줌
    if depth[a] > depth[b]:
        a, b = b, a
    diff = depth[b] - depth[a]
    for i in range(logN):
        if diff & 1 << i:
            b = dp[b][i]
    # 깊이 맞추고, 그 값이 같을 경우 최소 공통 조상은 a
    if a == b:
        print(a)
        continue
    # 루트 노드에서 부터 내려오면서 그 값이 달라지는 순간 찾기
    for i in range(logN - 1, -1, -1):
        if dp[a][i] != dp[b][i]:
            a = dp[a][i]
            b = dp[b][i]
    print(dp[b][0])

더 생각해 볼 것?

머리속의 생각을 좀더 알아보기 쉽게 설명하는 연습을 해야겠다.

코드나 내용 관련 조언, 부족한 점 및 질문 언제든 환영합니다!