[백준] 14942 - 개미 (Python)

개미집은 $1$번이 루트인 트리로 구조화할 수 있다.
모든 개미들이 에너지를 다 쓰기 전까지 루트로 올라가도록 구현해야 한다.

가장 간단하게 구현할 수 있는 것은 부모 노드로 한칸씩 올라가는 방식이다. 하지만 이렇게 되면 시간복잡도가 O(N^2)이 되기 때문에 비효율적이다. 다른 방법이 없을까?

한칸씩 올라가지 말고 2의 제곱수만큼 올라가면 $O(NlogN)$으로 구현이 가능하다.

먼저, 조상 노드를 저장하는 $parent$ 배열을 만든다. $parent[i][n]$ ($i$ => $2^i$ 번째 조상, $n$ => 현재 노드) 를 이용할 것이다.

$2^{i+1}$번째 조상은 $2^i$번째 조상의 $2^i$번째 조상이 된다. 아버지의 아버지는 할아버지인 것을 생각하면 이해가 쉽다. 점화식은 다음과 같다.
$parent[i+1][n] = parent[i][parent[i][n]]$

구현을 위해 루트부터 완전탐색을 통해 모든 노드에 대해 $parent$배열을 완성한다.

---

import sys
input = sys.stdin.readline

# parent 배열을 만드는 함수
def create_table(node):
	# 방문 체크
    visit[node] = 1 
    # 2^17 > 100000 이므로 i = 16까지만 저장
    for i in range(16):
        if parent[i][node]:
        	# anc: 조상, dist: 조상까지 거리
            anc, dist = parent[i][node]
            if parent[i][anc]:
            	# gnd_anc: 조상의 조상
                gnd_anc, gnd_dist = parent[i][anc]
                parent[i+1][node] = [gnd_anc, dist+gnd_dist]
    for b, c in nest[node]:
        if visit[b]:
            continue
        # 초기값 할당(첫 번째 조상)
        parent[0][b] = [node, c]
        create_table(b)

# node: 현재 노드, energy: 남은 에너지
def arrive(node, energy):
	# next_node: 다음으로 도착하는 조상
    next_node = node
    cost = 0
    for i in range(17):
        if parent[i][node]:
            anc_node, dist = parent[i][node]
            # energy가 허용하는 만큼 올라감
            if dist <= energy:
                next_node = anc_node
                cost = dist
            else:
                break
        else:
            break
    # 더 올라갈 수 없는 경우
    if next_node == node:
        return node
    return arrive(next_node, energy-cost)
    
n = int(input())
ant_energy = [0]
parent = [[0]*100001 for _ in range(17)]
nest = [[] for _ in range(n+1)]
visit = [0]*(n+1)
for _ in range(n):
    ant_energy.append(int(input()))
for _ in range(n-1):
    a, b, c = map(int, input().rstrip().split())
    nest[a].append((b, c))
    nest[b].append((a, c))
    
create_table(1)
for i in range(1, n+1):
    print(arrive(i, ant_energy[i]))