[백준/파이썬] 14942번 개미

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

14942번: 개미

 

---

문제 & 테스트 케이스

---

문제 해석

각 노드에 개미가 한마리 존재하며, 에너지가 주어진다.

개미들은 최대한 최상단 (1번 노드)로 이동해야 한다.

간선에는 이동하는데 필요한 에너지가 주어진다.

개미들이 최대한 1번노드 (최상단)으로 이동할 때 멈추게 되는 위치를 출력하는 문제이다.

---

코드

from collections import deque
input = __import__('sys').stdin.readline

n = int(input())
room = [0]
for _ in range(n):
    room.append(int(input()))

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


def dfs(power, room_number, start):
	# 이동이 완료되면 갱신
    ans[start] = room_number
    
    for togo, need in adj[room_number]:
    	# 다음 노드로 갈 수 있는 에너지를 가지고 있고 and 다음 노드가 더 상단에 위치한다면
        if power >= need and room_value[togo][1] < room_value[room_number][1]:
        	# 이동하기
            dfs(power - need, togo, start)

# bfs로 1번 룸에서 가장 가까운 방들 순서 정해주기.
def set_room_value_bfs():
	# 방 번호, 방의 level
    room_value = [(0, 0)] * (n + 1)
    value = 1
    q = deque()
    visit = [0] * (n + 1)
    visit[1] = 1
    q.append(1)
    while q:
        for _ in range(len(q)):
            now = q.popleft()
            room_value[now] = (now, value)
            for togo, need in adj[now]:
                if visit[togo] == 0:
                    visit[togo] = 1
                    q.append(togo)
        value += 1
    return room_value

room_value = set_room_value_bfs()
# print(room_value)
ans = [0] * (n + 1)

for i in range(1, n + 1):
    dfs(room[i], i, i)
    print(ans[i])

---

문제 풀이

bfs를 통해 이동 경로를 최적화시켜주고,

모든 노드들을 돌아가며 이동할 수 있는 최상단의 경로로만 이동시켜주어 답을 구했다.

 

1. bfs를 통해 room의 level 구하기.

주어진 예제를 그려보았다.

최상단의 위치한 노드로 부터 다음 레벨의 노드에 도달할 때 마다 level 의 값을 1씩 증가시켜주어야 한다.

즉, 1번 노드는 1레벨 / 2번 노드는 1레벨 / 3, 4번 노드는 3레벨이 되어야 한다.

이는 for문을 통해 단계별로 레벨을 구분하여 탐색하는 BFS를 통해 구별할 수 있다.

# bfs로 1번 룸에서 가장 가까운 방들 순서 정해주기.
def set_room_value_bfs():
	# 방 번호, 방의 level
    room_value = [(0, 0)] * (n + 1)
    value = 1
    q = deque()
    visit = [0] * (n + 1)
    visit[1] = 1
    q.append(1)
    while q:
        for _ in range(len(q)):
            now = q.popleft()
            room_value[now] = (now, value)
            for togo, need in adj[now]:
                if visit[togo] == 0:
                    visit[togo] = 1
                    q.append(togo)
        value += 1
    return room_value

따라서 room_value에 노드 번호, 노드 레벨을 담아두었다.

 

2. 조건에 따라 최상단으로 이동.

이제는 최적화된 경로로 이동할 수 있는 데이터를 알고 있으니, 조건에 따라 올라가기만 하면 된다.

두 가지 조건을 따라 올라가면 되는데,

먼저 내가 현재 가진 에너지가 충분한 경우 + 내 현재 노드의 레벨보다 더 낮은 노드의 레벨인 경우만

이동시켜주면 가장 최적으로 이동하게 되며 답을 갱신하게 된다.

def dfs(power, room_number, start):
	# 이동이 완료되면 갱신
    ans[start] = room_number
    
    for togo, need in adj[room_number]:
    	# 다음 노드로 갈 수 있는 에너지를 가지고 있고 and 다음 노드가 더 상단에 위치한다면
        if power >= need and room_value[togo][1] < room_value[room_number][1]:
        	# 이동하기
            dfs(power - need, togo, start)

조건문 실수..