[백준 1761번] 정점들의 거리 (Python3)

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

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

def findcoparent(x,y,c):
  global result
  if not c:
    if sparse[0][x] == sparse[0][y]:
      result += cost[0][x]+cost[0][y]
    else:
      result += cost[1][x]+cost[1][y]
    return
  if sparse[c][x] == sparse[c][y]:
    findcoparent(x,y,c-1)
  else:
    result += cost[c][x]+cost[c][y]
    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,w = map(int,input().split())
  graph[x-1].append((y-1,w))
  graph[y-1].append((x-1,w))

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

dq = deque([(0,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 for i in range(logd)]
cost = [cost] + [[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]]
    cost[i][n] = cost[i-1][n]+cost[i-1][sparse[i-1][n]]

M = int(input())
for _ in range(M):
  x,y = map(lambda x:int(x)-1,input().split())
  result = 0
  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:
    findcoparent(x,y,int(log2(depth[x]-1)))
  print(result)

LCA [백준 11438번] LCA 2 (Python3) (tistory.com) 알고리즘 기본문제.

 

cost에 대한 sparse table을 만들어주고 LCA 알고리즘으로 해결한다.

LCA 알고리즘을 알면 매우 쉽게 해결할 수 있는 문제