문제

민오는 1번부터 N번까지 총 N개의 문제로 되어 있는 문제집을 풀려고 한다. 문제는 난이도 순서로 출제되어 있다. 즉 1번 문제가 가장 쉬운 문제이고 N번 문제가 가장 어려운 문제가 된다.

어떤 문제부터 풀까 고민하면서 문제를 훑어보던 민오는, 몇몇 문제들 사이에는 '먼저 푸는 것이 좋은 문제'가 있다는 것을 알게 되었다. 예를 들어 1번 문제를 풀고 나면 4번 문제가 쉽게 풀린다거나 하는 식이다. 민오는 다음의 세 가지 조건에 따라 문제를 풀 순서를 정하기로 하였다.

1. N개의 문제는 모두 풀어야 한다.
2. 먼저 푸는 것이 좋은 문제가 있는 문제는, 먼저 푸는 것이 좋은 문제를 반드시 먼저 풀어야 한다.
3. 가능하면 쉬운 문제부터 풀어야 한다.

예를 들어서 네 개의 문제가 있다고 하자. 4번 문제는 2번 문제보다 먼저 푸는 것이 좋고, 3번 문제는 1번 문제보다 먼저 푸는 것이 좋다고 하자. 만일 4-3-2-1의 순서로 문제를 풀게 되면 조건 1과 조건 2를 만족한다. 하지만 조건 3을 만족하지 않는다. 4보다 3을 충분히 먼저 풀 수 있기 때문이다. 따라서 조건 3을 만족하는 문제를 풀 순서는 3-1-4-2가 된다.

문제의 개수와 먼저 푸는 것이 좋은 문제에 대한 정보가 주어졌을 때, 주어진 조건을 만족하면서 민오가 풀 문제의 순서를 결정해 주는 프로그램을 작성하시오.

풀이

우선 선행 문제를 풀어야 후행 문제를 풀 수 있다는 점에서 위상정렬 알고리즘을 사용해야 한다는 것을 알 수 있다. 그러나 위상정렬 문제에서 자주 보이는 "답이 여러 개인 경우에는 아무거나 출력한다."라는 조건이 없는 것을 보아 정답은 유일성을 가짐을 확인할 수 있다.

보통 위상정렬 알고리즘은 다음과 같이 사용된다.

1. 진입 차수 (해당 노드의 선행 노드의 개수)가 0인 노드를 queue에 넣는다.
2. 해당 노드로부터 갈 수 있는 다음 노드들에 대한 연결을 끊는다.
3. 1~2번 과정을 반복한다.

그러나 본 문제에는 "가능한 쉬운 문제부터 풀어야한다"라는 조건이 주어진다.

가령, 제일 처음에 1, 3, 4번 문제가 선행문제가 존재하지 않고, 2번 문제의 선행문제는 1번이라고 가정해보자. 위 알고리즘대로라면 queue에 1번 3번 4번이 삽입된 후에 2번이 삽입된다.

그러나 1번을 이미 푼 상태라면, 3번 4번을 풀 필요 없이 2번을 풀 수 있으며, 2번을 먼저 푸는 것이 "가능한 쉬운 문제부터 풀어야한다"라는 조건을 만족한다.

따라서 queue 대신에 heap을 사용하는 것이 본 문제의 핵심이라고 할 수 있다.

코드

import sys
input = sys.stdin.readline
import heapq

# get inputs
num_qs, num_info = map(int, input().split()) # 문제의 개수, 먼저 풀면 좋은 정보의 개수

nums_pre = [0] * (num_qs+1) # nums_pre[i] = i번 문제 이전에 풀어야 하는 문제의 개수
q_nexts = [[] for _ in range(num_qs +1)] # q_nexts[i] = i번 문제 이후에 풀 수 있는 문제들

# 문제 선행관계에 대한 정보 업데이트
for _ in range(num_info):
    q_pre, q_next = map(int, input().split())
    q_nexts[q_pre].append(q_next)
    nums_pre[q_next] += 1

# 위상정렬
heap = []
path = []

# 선행 문제가 없는 문제들 heap에 추가 
for i in range(1, num_qs +1):
    if nums_pre[i] == 0:
        heapq.heappush(heap, i)
        nums_pre[i] -= 1

while heap:

    question = heapq.heappop(heap)
    path.append(question)
    
    # 후행 문제에 대한 link 끊기
    for q_next in q_nexts[question]:
        nums_pre[q_next] -= 1

        # 새로이 선행 문제 없는 문제 찾기
        if nums_pre[q_next] == 0:
            heapq.heappush(heap, q_next)

print(*path)

비고 (틀린 풀이)

while heap:

    question = heapq.heappop(heap)
    path.append(question)
    
    # 후행 문제에 대한 link 끊기
    for q_next in q_nexts[question]:
        nums_pre[q_next] -= 1

    # 새로이 선행 문제 없는 문제 찾기
    for i in range(1, num_qs +1):
        if nums_pre[i] == 0:
            heapq.heappush(heap, i)
            nums_pre[i] = -1

처음에는 후행 문제에 대해 link를 끊은 다음에, 1번 문제부터 n번 문제까지 전부 탐색하면서 진입차수가 0인 문제를 검색하였다.

그러나 이는 매우 비효율적인 탐색이기 때문에 시간초과가 발생하였다.

가령, $k$번 문제가 $k+1, k+2, k+3$번 문제의 선행문제라고 가정해보자.
새로이 $k$번 문제를 푼다면, $k+1, k+2, k+3$번 문제의 진입 차수만 업데이트 되기 때문에 굳이 1번 문제부터 $n$번 문제까지 탐색할 필요가 없다. 새로이 업데이트 된 $k+1, k+2, k+3$번 문제만 탐색하면 되는 일이다.

지금까지 그것도 모르고 위상정렬 문제는 모두 pypy로만 풀어왔었는데 (pypy로 돌리면 $n$개의 문제를 모두 탐색해도 시간이 초과되지 않는다.), 내 위상정렬 코드의 문제를 알 수 있는 좋은 문제였다.