[골드 2] 백준 2637 - 장난감 조립 (파이썬)

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

풀이

• 어떤 완제품을 만들기 위해 여러 부품이 필요하고, 그 부품들 또한 다른 부품들로 조립되는 구조
• 완제품 번호는 항상 N번이고, 기본 부품은 다른 부품으로 구성되지 않는 것
• 기본 부품마다 완제품을 1개 조립할 때 필요한 개수를 구하는 문제

입력 처리 및 그래프 구성은 아래와 같이 구현한다.

adj = [[] for _ in range(N)]   # 인접 리스트 (X -> (Y, K): X는 Y를 K개 사용함)
count = [0] * N                # 진입 차수 저장

for i in range(M):
    X, Y, K = map(int, input().split())
    X -= 1
    Y -= 1
    adj[X].append((Y, K))      # X를 만들기 위해 Y가 K개 필요하다
    count[Y] += 1              # Y의 진입 차수 증가

이 구조는 역방향으로써, 즉 X를 만들기 위해 Y가 필요하면 X → Y 방향 간선을 생성한다.

진입 차수는 Y에 대해 증가한다.

위상 정렬을 시작하기 위해서 초기 세팅을 아래와 같이 진행한다.

answer = [0] * N
answer[N - 1] = 1

완제품 번호는 N, 즉 N - 1 인덱스이다.

완제품 1개를 만든다고 가정하고 시작한다. (answer[N - 1] = 1)

for i in range(N):
    if count[i] == 0:
        queue.append(i)

진입 차수가 0인 노드를 큐에 넣는다. 여기서 기본 부품들이 먼저 들어간다.

while queue:
    u = queue.popleft()

    for v, w in adj[u]:
        answer[v] += w * answer[u]
        count[v] -= 1
        if count[v] == 0:
            queue.append(v)

u는 현재 처리중인 부품이다.

u를 통해서 만들 수 있는 상위 부품 v에 대해서, answer[v] += w * answer[u]는 u가 answer[u]개 필요하므로, v에 필요한 u의 양을 누적한다.

진입 차수가 0이 되면 큐에 삽입하여, 위상 정렬을 진행한다.

for i in range(N):
    if len(adj[i]) == 0:      # 기본 부품 (다른 부품을 필요로 하지 않음)
        print(i + 1, answer[i])

최종적으로 adj[i]가 비어있는 경우는 기본 부품이다.

해당 부품이 얼마나 필요한지 answer[i]에 저장되어 있다.

코드

# 백준 2637 - 장난감 조립
# 분류 : 위상정렬

from collections import deque
N = int(input())
M = int(input())

adj = [[] for _ in range(N)]
count = [0] * N

for i in range(M) :
    X, Y, K = map(int, input().split())
    X -= 1
    Y -= 1

    adj[X].append((Y, K))
    count[Y] += 1

answer = [0] * N
answer[N - 1] = 1

queue = deque()
for i in range(N) :
    if count[i] == 0 :
        queue.append(i)

while len(queue) != 0 :
    u = queue.popleft()

    for v, w in adj[u] :
        answer[v] += w * answer[u]
        count[v] -= 1

        if count[v] == 0 :
            queue.append(v)

for i in range(N) :
    if len(adj[i]) == 0 :
        print(i + 1, answer[i])