[골드 5] 백준 14567 - 선수과목 (Prerequisite) (파이썬)

분류 : 다이나믹 프로그래밍

링크 : https://www.acmicpc.net/problem/14567

풀이

총 N개의 과목이 주어지며, M개의 선수 조건이 주어진다.

모든 과목에 대해 가장 빠른 이수 학기를 출력해야한다.

N, M = map(int, input().split())
adj = [[] for _ in range(N)]

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

adj[b].append(a)는 b를 듣기 위해 a를 먼저 들어야 한다는 조건이므로, b에서 a로 역방향 간선을 저장한다.

즉, b의 선수 과목들을 adj[b]에 저장한다.

cache = [-1] * N

cache[i]는 i번 과목을 이수할 수 있는 가장 빠른 학기를 저장한다.

def dp(u) :
    if cache[u] != -1 :
        return cache[u]

    if len(adj[u]) == 0 :
        ret = 1
    else :
        ret = max([dp(v) + 1 for v in adj[u]])
    cache[u] = ret
    return ret

dp(u)는 과목 u를 듣기 위한 최소 학기 수를 반환한다.

선수 과목이 없으면 바로 1학기에 들을 수 있으며, 선수 과목이 있다면 그 중 가장 늦게 들을 수 있는 과목 + 1학기가 돼야 한다.

메모이제이션으로 중복 계산을 방지한다.

코드

# 백준 14567 - 선수과목 (Prerequisite)
# 분류 : 다이나믹 프로그래밍

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 6)

N, M = map(int, input().split())
adj = [[] for _ in range(N)]

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

    adj[b].append(a)

cache = [-1] * N
def dp(u) :
    if cache[u] != -1 :
        return cache[u]

    if len(adj[u]) == 0 :
        ret = 1
    else :
        ret = max([dp(v) + 1 for v in adj[u]])
    cache[u] = ret
    return ret

for i in range(N) :
    print(dp(i), end = " ")