[백준 11438번] LCA 2 (Python3)

import sys
input = sys.stdin.readline
from collections import deque
from math import log2

def findparent(x,d):
  for i in range(logd+1):
    if d&(1<<i):
      x = sparse[i][x]
  return x

def findcoparent(x,y,c):
  if not c:
    if parent[x] == parent[y]:
      return parent[x]
    return sparse[1][x]
  if sparse[c][x] == sparse[c][y]:
    return findcoparent(x,y,c-1)
  return findcoparent(sparse[c][x],sparse[c][y],c-1)

N = int(input())

graph = [[] for i in range(N)]
for _ in range(N-1):
  x,y = map(lambda x:x-1,map(int,input().split()))
  graph[x].append(y)
  graph[y].append(x)

parent = [0]*N
depth = [0]*N

dq = deque([(0,1)])
while dq:
  now,d = dq.popleft()
  depth[now] = d
  for next in graph[now]:
    if depth[next]:
      continue
    parent[next] = now
    dq.append((next,d+1))

logd = int(log2(d))
sparse = [parent]+[[0]*N for i in range(logd)]
for i in range(1,logd+1):
  for n in range(N):
    sparse[i][n] = sparse[i-1][sparse[i-1][n]]

M = int(input())
for _ in range(M):
  x,y = map(lambda x:x-1,map(int,input().split()))
  if depth[x] > depth[y]:
    x = findparent(x,depth[x]-depth[y])
  elif depth[x] < depth[y]:
    y = findparent(y,depth[y]-depth[x])
  if x == y:
    print(x+1)
  else:
    print(findcoparent(x,y,int(log2(depth[x]-1)))+1)

LCA [백준 11437번] LCA (Python3) (tistory.com) 문제를 sparse table [백준 17435번] 합성함수와 쿼리 (Python3) (tistory.com)를 이용해서 효율적으로 해결하는 문제이다.

 

sparse table을 이용해야 한다는 사실을 몰랐으면 아마 절대 풀지 못했을 것 같다.

 

기본적인 알고리즘은 LCA1과 똑같다.

1. 1번노드를 기준으로 트리를 작성, 모든 노드의 부모노드와 깊이를 저장한다.

2. 두 노드의 깊이가 다르면 같은 깊이까지 올라온다.

3. 현재의 깊이에서 부모가 같아질 때까지 위로 올라온다. 

 

이때 2,3번 과정을 sparse table을 이용해서 O(N)에 끝나는 작업을 O(logN)으로 해결하는 것이 바로 LCA2의 핵심이다.

 

 

오늘의 교훈) LCA 알고리즘과 sparse table을 적절하게 활용하자