[Python] 백준 / gold / 1766 : 문제집

문제 링크 : https://www.acmicpc.net/problem/1766

위상 정렬 문제다. 이 문제를 풀고 이코테 그래프 이론 부분을 복습해야겠다고 느꼈다. 처음엔 먼저 풀어야 하는 번호를 딕셔너리에 담아 힙 자료구조로 저장했다. 딕셔너리에 먼저 풀어야 할게 있다면 DFS 를 진행하는 방법으로 풀었는데, 오답처리가 났다.

테스트 케이스를 만들어서 여러 개 시도해봤는데 반례를 찾기 힘들었다. 반례는 이런 경우였다. DFS 로 무조건 깊게 들어가면 안되고, 들어갈 때마다 먼저 풀 수 있는 번호가 있는지 체크를 해줬어야 됐다.

이렇게 푸는건 너무 쓸데없이 복잡하기 때문에, 이 문제는 위상정렬 + 힙으로 푸는게 정석이었다.

기억할 것 : 위상정렬은 inDegree 를 이용해서 푼다.

오답 코드

from collections import defaultdict
import sys
import heapq
N, M = map(int, sys.stdin.readline().split())
priority = defaultdict(list)
visit = [False]*(N+1)
answer = []

for _ in range(M):
    A, B = map(int, sys.stdin.readline().split())
    heapq.heappush(priority[B], A)


def dfs(x):
    if x not in priority:
        answer.append(x)
        visit[x] = True
        return

    h = priority[x]
    while h:
        now = heapq.heappop(h)
        if visit[now]:
            continue
        dfs(now)

    answer.append(x)
    visit[x] = True
    return


for n in range(1, N+1):
    if visit[n]:
        continue
    # 우선순위 사전에 없으면
    if n not in priority:
        answer.append(n)
        visit[n] = True
    # 우선순위 사전에 있으면
    else:
        dfs(n)

for a in answer:
    print(a, end=' ')

반례

6 7
5 6
5 2
2 4
4 3
2 1
6 1
1 3

정답 : 5 2 4 6 1 3

output : 5 2 6 1 4 3

정답 코드

import sys
import heapq

N, M = map(int, sys.stdin.readline().split())
graph = [[] for _ in range(N+1)]
inDegree = [0]*(N+1)
answer = []

for _ in range(M):
    A, B = map(int, sys.stdin.readline().split())
    graph[A].append(B)
    inDegree[B] += 1

h = []
for i in range(1, N+1):
    if inDegree[i] == 0:
        heapq.heappush(h, i)

while h:
    now = heapq.heappop(h)
    answer.append(now)
    for neighbor in graph[now]:
        inDegree[neighbor] -= 1
        if inDegree[neighbor] == 0:
            heapq.heappush(h, neighbor)

for x in answer:
    print(x, end=' ')