[백준 13511번] 트리와 쿼리 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 calcost(x,d):
  SUM = 0
  for i in range(logd+1):
    if d&(1<<i):
      SUM += sumcost[i][x]
      x = sparse[i][x]
  return SUM

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+1)]
for i in range(N-1):
  x,y,w = map(int,input().split())
  graph[x].append((y,w))
  graph[y].append((x,w))

parent = [1]*(N+1)
depth = [0]*(N+1)
cost = [0]*(N+1)

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

logd = int(log2(d))
sparse = [parent] + [[0]*(N+1) for i in range(logd)]
sumcost = [cost]+[[0]*(N+1) for i in range(logd)]

for i in range(1,logd+1):
  for n in range(1,N+1):
    sparse[i][n] = sparse[i-1][sparse[i-1][n]]
    sumcost[i][n] = sumcost[i-1][n]+sumcost[i-1][sparse[i-1][n]]

M = int(input())
for _ in range(M):
  Q = [*map(int,input().split())]
  u,v = Q[1],Q[2]
  
  if depth[u]>depth[v]:
    u = findparent(u,depth[u]-depth[v])
  elif depth[v]>depth[u]:
    v = findparent(v,depth[v]-depth[u])
  if u == v:
    p = u
  else:
    p = findcoparent(u,v,int(log2(depth[u]-1)))

  u,v = Q[1],Q[2]
  if Q[0] == 1:
    print(calcost(u,depth[u]-depth[p])+calcost(v,depth[v]-depth[p]))
  else:
    k = Q[3]-1
    if depth[u]-depth[p] == k:
      print(p)
    elif depth[u]-depth[p] > k:
      print(findparent(u,k))
    else:
      print(findparent(v,depth[v]+depth[u]-2*depth[p]-k))

LCA [백준 11438번] LCA 2 (Python3) (tistory.com) 와 sparse table을 이용하는 또 다른 문제이다.

LCA 알고리즘을 알고 있다는 가정하에 쉽게 해결 가능하다.

 

알고리즘을 요약하면,

1. sparse table을 만들 때, cost의 합에 대한 sparse table도 같이 만들어준다.

2. LCA 알고리즘을 사용하여 u,v의 공통조상 p를 찾는다.

3. 1번 쿼리의 경우, u~p까지의 cost와 v~p까지의 cost의 합을 출력한다.

4. 2번 쿼리의 경우, u와 p의 높이차이보다 k가 작은 경우 u의 k번째 조상을 출력한다. 아닐 경우, v의 (u와 p의 깊이차이+v와 p의 깊이차이-k) 번째 조상을 출력한다.

 

 

오늘의 교훈) LCA 알고리즘은 유용하다.