[Python] 백준 1005번: ACM Craft

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

1005번: ACM Craft

 

import sys
input = sys.stdin.readline


def dfs(building, depth):
    global buildingsOnDepth

    if buildingsOnDepth.get(depth):
        if building in buildingsOnDepth[depth]:
            return

        buildingsOnDepth[depth].add(building)
    else:
        buildingsOnDepth[depth] = set([building])

    if building == targetBuilding:
        return

    for next in nextBuildings[building]:
        dfs(next, depth+1)


def searchBuilding(building):
    global costsToBuild

    if building == targetBuilding:
        return

    for next in nextBuildings[building]:
        costsToBuild[next] = max(
            costsToBuild[next], costsToBuild[building] + costs[next])


T = int(input())
for _ in range(T):
    N, K = map(int, input().split())

    costs = [0]
    costs += list(map(int, input().split()))

    nextBuildings = {i: [] for i in range(N+1)}
    requiredBuildings = {i: [] for i in range(N+1)}

    for _ in range(K):
        cur, next = map(int, input().split())
        nextBuildings[cur].append(next)
        requiredBuildings[next].append(cur)

    # 헤더 빌딩(0번)에, 처음부터 지을 수 있는 건물들 연결
    for b in range(1, N+1):
        if requiredBuildings[b] == []:
            nextBuildings[0].append(b)
            requiredBuildings[b].append(0)

    targetBuilding = int(input())

    # 깊이 별로 빌딩들을 나눔
    buildingsOnDepth = {}
    dfs(0, 0)

    # 깊이에 따른 탐색
    costsToBuild = [0 for _ in range(N+1)]
    for buildings in buildingsOnDepth.values():
        for b in buildings:
            searchBuilding(b)

    print(costsToBuild[targetBuilding])

 

우선 기본 건물(다른 건물을 필요로 하지 않는 건물)이 여러 개일 수도 있으며, 무조건 1번 건물이 기본 건물이라고 하지 않았다. 따라서 나는 기본 건물들을 '0번 건물'에 묶어서 마치 연결리스트의 헤더와 같은 기능을 만들고자 했다.

이후에 0번 건물을 필두로, dfs를 수행하면서 '깊이'에 따른 건물들을 묶기로 하였다. 여기서 buildingOnDepth라는 리스트가 쓰이는데, 이 리스트의 원소에 각각의 깊이에서의 빌딩들을 집합으로 저장한다.

이렇게 buildingOnDepth를 초기화했으면, 마지막으로 깊이 순서로 빌딩들을 탐색해준다. 이 탐색에서 costsToBuild라는 리스트를 초기화해주는데, costsToBuild[i]는 i 빌딩을 짓는 데 걸리는 시간이다.