[Python][백준] 14942번 개미

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

공부 날짜 : 2023.02.10
정답 참조 여부 : X

트리구조로된 개미 굴에서 각 방에 있는 개미들이 축척된 에너지를 이용해서 출구까지 나오려고 한다.

간선의 가중치 만큰 축적된 에너지를 소모한다고 할 때 최대한 도달 할 수 있는 방의 위치를 구하는 문제이다.

---

트리를 dfs로 탐색하며

탐색하지 않았으면 방문 처리를 하고 다음 노드를 탐색 스택에 넣기

탐색은 했지만, 자식노드를 모두 탐색한게 아니라면 이전방에서 넘어온 개미를 이번방에 추가하기

자식노드를 모두 탐색했다면 지금 방에 있는 개미들 모두에게 다음방으로 넘어가는 비용만큼 에너지를 빼고 에너지가 남아있는 개미만 다음방으로 넘어가기

로 로직을 구성했다.

---

일단 이렇게 구성을 했지만, 구현이 만만치 않았다.

자식노드를 모두 탐색했는지 판단하기위해 자식노드를 넣기전에 자신의 노드를 음수로 탐색 스택에 넣었고, 매 자식 노드를 넣고 자신의 노드를 추가로 입력했다.

그렇게 자식노드에서 탐색을 마치고 자신의 노드로 돌아와 개미들이 들어오기 때문이다.

또한, 개미가 다음방으로 넘어가는지 판단한 후에 부모 노드로 돌아가는 것도 좀 어려웠는데
다음 방으로 넘어가는 개미들의 스택을 만들었고 매번 자신의 노드 탐색이 끝나면 부모 노드로 넘어가기 때문에(부모노드에서 자식노드를 입력하기 전에 자신의 노드를 입력해줬기 때문) 부모 노드에서 탐색이 완료되었으면 이번방으로 넘어오는 개미들을 전부 받아 들였다.(방에 개미를 넣으면서 개미들의 위치도 옮긴다.)

마지막으로 루트노드인 1로 돌아오면 탐색을 종료하고, 결과를 출력했다.

---

솔직히 풀면서 많이 어려웠지만, 설명할건 딱히 없다.... 로직 그대로 구현만 했기 때문
자잘한 오타로 어려움은 겪었지만 로직을 그대로 구현했더니 바로 정답으로 나왔다.

소스코드

import sys
input = sys.stdin.readline
###############################################
n = int(input())

ant_energy = [0]
room = [[i] for i in range(n+1)]
ant_position = [i for i in range(n+1)]

for _ in range(n):
    ant_energy.append(int(input()))


tunnel = [[] for _ in range(n+1)]

for _ in range(n-1):
    a, b, length = map(int, input().split())

    tunnel[a].append((b, length,))
    tunnel[b].append((a, length,))

# print(tunnel)

visit = [1]
visited = [False for _ in range(n+1)]
cost_stack = []

go_next_ant = []

while visit:
    now_room = visit.pop()

    # 음수가 저장되어 있다면 자식노드를 모두 탐색하고 돌아온것
    comeback_check = False
    if now_room < 0:
        comeback_check = True
        now_room = -now_room

	# 루트노드로 돌아왔는데, 루트노드를 방문했다면 탐색 끝
    if now_room == 1 and visited[1]:
        while go_next_ant:
            room[now_room].append(go_next_ant[-1])
            ant_position[go_next_ant.pop()] = now_room
        continue


    # 이전에 방문했다면
    if visited[now_room]:
        # 이번 방으로 넘어오는 개미들은 전부 추가하고 개미들 옮기기
        while go_next_ant:
            room[now_room].append(go_next_ant[-1])
            ant_position[go_next_ant.pop()] = now_room

        # 자식노드를 모두 돌고 왔다면
        if comeback_check:
            # 이번방에서 부모방으로 가는데 필요한 비용
            cost = cost_stack.pop()

            # 지금 방에 있는 모든 개미들에 대해서
            for i in room[now_room]:
                # 다음방으로 넘어가기 위한 비용 빼기
                ant_energy[i] -= cost
                # 넘어가도 에너지가 남아있으면 다음방 가는 리스트에 추가가
                if ant_energy[i] >= 0:
                    go_next_ant.append(i)

        continue

    visited[now_room] = True

    # 자식노드 모두 탐색 후 자신에게 돌아오기 위한 추가
    visit.append(-now_room)

    for room_, cost in tunnel[now_room]:
        if visited[room_]:
            continue

        visit.append(room_)
        visit.append(now_room)
        cost_stack.append(cost)

    # visit.pop()


for i in range(1, n+1):
    print(ant_position[i])