[백준 15480번] LCA와 쿼리 (Python3)

import sys
input = sys.stdin.readline
from collections import *

def findparent(x,d):
  for i in range(l+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)

def LCA(x,y):
  dx,dy = depth[x],depth[y]
  if dx>dy:
    x = findparent(x,dx-dy)
  elif dx<dy:
    y = findparent(y,dy-dx)
  if x==y:
    return x
  return findcoparent(x,y,l)

N = int(input())

graph = [[] for i in range(N+1)]
for _ in range(N-1):
  a,b = map(int,input().split())
  graph[a].append(b); graph[b].append(a)

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

l = len(bin(d))-3
sparse = [parent] + [[0]*(N+1) for i in range(l)]
for i in range(1,l+1):
  for n in range(1,N+1):
    sparse[i][n] = sparse[i-1][sparse[i-1][n]]

M = int(input())
for _ in range(M):
  r,a,b = map(int,input().split())
  lca,lca1,lca2 = LCA(a,b),LCA(r,a),LCA(r,b)
  if lca1==lca2:
    print(lca)
  elif lca1!=lca:
    print(lca1)
  else:
    print(lca2)

간단한 LCA [백준 11438번] LCA 2 (Python3) (tistory.com) 알고리즘 응용문제

sparse table을 이용한 LCA 최적화 알고리즘을 알고 있다면 쉽게 해결할 수 있다.

 

풀이를 설명하면,

1. LCA(a,b), LCA(r,a), LCA(r,b)를 구한다.

2. LCA(r,a)와 LCA(r,b)가 같은 경우, LCA(a,b)를 출력한다.

3. 2가 아닌 경우, LCA(r,a), LCA(r,b) 중에서 LCA(a,b)와 다른 값을 출력한다. 

 

2가 성립하는 이유는 LCA(r,a), LCA(r,b)가 같으면 r은 a-b 경로의 바깥에 존재한다는 뜻이 된다. 따라서 LCA(a,b)는 루트가 r로 바뀌어도 그대로이다.

3이 성립하는 이유는 예를들어 만약 LCA(r,a) != LCA(a,b)라면, r은 a-b 경로상에서 LCA(a,b)를 기준으로 a쪽에 가깝게 위치하고 있다는 뜻이 된다. 따라서 루트가 r로 바뀌면 LCA(a,b)는 LCA(r,a)가 된다.

 

r이 a-b 경로 내부에 존재하는 경우와 그렇지 않은 경우를 고려해주면 쉽게 해결할 수 있다.

 

 

오늘의 교훈) 트리의 특수성을 잘 활용하자