[백준 14942번] 개미 (Python3)

import sys,collections,math
input = sys.stdin.readline

def climb(x,e,c):
  if not c:
    if cost[0][x] <= e:
      return sparse[0][x]
    return x
  if cost[c][x] <= e:
    return climb(sparse[c][x],e-cost[c][x],c-1)
  return climb(x,e,c-1)

N = int(input())
energy = [0]+[int(input()) for i in range(N)]

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

parent,cost,depth = [[0]*(N+1) for i in range(3)]
parent[1] = 1
dq = collections.deque([(1,1)])
while dq:
  now,d = dq.popleft()
  depth[now] = d
  for next,c in graph[now]:
    if depth[next]:
      continue
    parent[next],cost[next] = now,c
    dq.append((next,d+1))

logd = int(math.log2(d))
sparse = [parent]+[[0]*(N+1) for i in range(logd)]
cost = [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]]
    cost[i][n] = cost[i-1][n]+cost[i-1][sparse[i-1][n]]

for i in range(1,N+1):
  print(climb(i,energy[i],int(math.log2(depth[i]))))

기본적인 sparse table 문제이다.

 

어차피 위로만 올라가기 때문에 최소공통조상을 찾을 필요는 없고, 부모에 대한 sparse table과 cost에 대한 sparse table을 만들어주면 된다.

 

 

오늘의 교훈) 트리문제는 항상 sparse table 염두하기