[백준/14942] 개미 (Python)

Problem : https://www.acmicpc.net/problem/14942

14942번: 개미

 

Difficulty : Platinum 5

 

Status : Solved

 

Time : 00:16:15

 

---

 

문제 설명

 

 

개미집은 n개의 방으로 구성되어 있으며 n개의 방은 1번부터 n번 까지 번호가 부여되어 있다. 그 중에서 1번 방은 지면에 바로 연결되어 있는 방이다. 각 방들은 서로 굴을 통해 연결되어 있다. 각 굴을 이동하기 위해서는 굴의 길이만큼 에너지가 소모된다.

개미는 집짓기의 달인이기 때문에 불필요한 굴은 짓지 않는다. 그래서 굴을 타고 한 방에서 다른 방으로 갈 수 있는 경로는 항상 존재하며 유일하다. 임의의 두 방 사이의 거리는 두 개의 방을 연결하는 경로를 구성하는 굴의 길이의 합이다.

겨울잠을 자던 개미들은 겨울잠에서 깨어나 지면으로 올라가 햇살을 보고 싶어한다. 그렇기 때문에 지면과 연결된 1번 방으로 이동을 하려고 한다. 하지만 불행하게도 개미는 긴 겨울잠을 자느라 축적해 놓은 에너지가 적다. 그래서 개미는 에너지를 1번 방에 도달하기 전에 모두 소모 할 수도 있다. 이렇게 에너지가 0이 된 개미는 더 이상 움직일 수 없다. 또한 1번 방에 도착한 개미는 더 이상 움직이지 않는다.

현재 모든 방에는 개미가 한 마리씩 있고 각각의 개미는 각자 축적된 에너지를 가지고 있다. 잠에서 깨어난 모든 개미는 1번 방을 향해서 이동한다. 이때 각각의 개미에 대해 도달할 수 있는 방 중에서 가장 1번 방에 가까운 방의 번호를 출력하시오.

 

 

입력 및 출력

 

입력

 

자연수 n이 주어진다. n은 방의 개수이다. (1 ≤ n ≤ 10^5) 다음 n개의 줄에는 차례대로 현재 각각의 개미가 보유하고 있는 에너지 값이 주어진다. i+1번째 줄에는 i번째 방에 있는 개미가 가진 에너지를 나타내는 100,000이하의 자연수 값이 주어진다. 이후 n-1개의 줄에는 두 개의 방을 연결하는 굴의 정보가 3개의 정수 a b c 으로 주어진다. a, b는 연결된 방을 의미하고 c는 이 굴의 길이를 의미한다. 굴의 길이는 10,000 이하의 자연수이다.

 

출력

 

n개의 줄을 출력한다. i번째 줄에는 i번 방에 있던 개미가 도달할 수 있는 방 중에 1번 방과 가장 가까운 방의 번호를 출력한다.

 

입력 예시

 

4
10
8
22
18
1 2 10
2 3 10
2 4 10

 

 

출력 예시

 

1
2
1
2

 

 

---

 

풀이

 

 

개미마다 올라갈 수 있는 최대 방의 번호를 구하는 게 문제이지만, 단순히 DFS로 풀면 개미 한마리당 O(N)의 시간복잡도를 가지게 된다. 따라서 총 시간복잡도가 O(N^2)이 되므로, 이 방법을 곧이곧대로 쓸 수는 없을 것 같다.

 

최소 공통 조상 문제들을 다루며, 우리는 이런 식으로 부모 노드들을 재귀적으로 탐색할 때 희소 배열을 사용하는 방식을 채용했다. 즉 (1번쩨, 2번째, 4번째, .... 2^n번째) 조상 노드들과, 그 노드들로 한 번에 가기 위해 필요한 비용을 기록해 둔다.

 

d번째 조상 노드들을 업데이트할 때, parent[node][d-1] = (pnode, pcost)라고 하자. 2^(d-1)번째 조상의 노드 번호가 pnode, 그리고 그 비용이 pcost가 된다. 이 때, parent[pnode][d-1] = (gpnode, gpcost)가 되며, 우리가 구하고자 하는 parents[node][d] = (gpnode, pcost + gpcost)를 만족하게 된다. 초기값인 parent[node][d]는 루트 노드(1번)로부터 그래프 탐색을 통해 구할 수 있다.

 

후에 탐색을 진행할 때, 2^n번째부터 1번째 조상까지중 현재 비용으로 갈 수 있는 가장 높은 조상으로 그리디하게 이동하는 게 포인트가 되겠다. 재귀적으로 현재 남은 에너지로 갈 수 있는 최대 조상 노드로 한 번에 이동하면 되는 식이다. N = 10^5이므로 최대 logN 이 20보다 작게 된다. 즉 어떤 노드에서라도 약 20번의 횟수 내에 조상 노드까지 최대한 빠르게 도달할 수 있다. 이 때 총 시간복잡도는 O(NlogN)이 된다.

 

 

풀이 코드

from collections import defaultdict
import sys
input = sys.stdin.readline
D = 20

N = int(input())

ants = [int(input()) for _ in range(N)]
edge_dict = defaultdict(list)

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

visited = [False]*N
visited[0] = True
parents = [[(0, 0) for _ in range(D)] for _ in range(N)]
result = [0]*N

def dfs(node) :
  for nxt, cost in edge_dict[node] :
    if visited[nxt] :
      continue
    visited[nxt] = True
    parents[nxt][0] = (node, cost)
    dfs(nxt)

def parent_init() :
  for d in range(1, D) :
    for node in range(N) :
      pnode, pcost = parents[node][d-1]
      gpnode, gpcost = parents[pnode][d-1]
      parents[node][d] = (gpnode, pcost+gpcost)

def search(idx, cur, left) :
  if cur == 0 :
    result[idx] = cur + 1
    return
  for d in range(D-1, -1, -1) :
    pnode, pcost = parents[cur][d]
    if left >= pcost :
      search(idx, pnode, left-pcost)
      return
  result[idx] = cur + 1

dfs(0)
parent_init()
for i in range(N) :
  search(i, i, ants[i])
print(*result, sep = '\n')

풀이 완료!