[Python] 백준 14938번: 서강그라운드

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

14938번: 서강그라운드

 

import sys
import heapq
input = sys.stdin.readline


def getPossibleItemCount(start):
    dp = [float('inf')]*(n+1)
    dp[start] = 0
    todo = [(0, start)]

    while todo:
        curDist, curNode = heapq.heappop(todo)

        if curDist > dp[curNode]:
            continue

        for next, dist in graph[curNode]:
            nextDist = curDist + dist

            if nextDist <= m and nextDist < dp[next]:
                dp[next] = nextDist
                heapq.heappush(todo, (nextDist, next))

    possibleItemCount = 0

    for i in range(1, n+1):
        if dp[i] != float('inf'):
            possibleItemCount += itemCounts[i]

    return possibleItemCount


if __name__ == '__main__':
    # 입력
    n, m, r = map(int, input().split())
    itemCounts = [-1] + list(map(int, input().split()))

    graph = [[] for _ in range(n+1)]
    for _ in range(r):
        a, b, l = map(int, input().split())
        graph[a].append((b, l))
        graph[b].append((a, l))

    # 처리
    maxItemCount = 0
    for i in range(1, n+1):
        maxItemCount = max(maxItemCount, getPossibleItemCount(i))

    # 출력
    print(maxItemCount)

 

다익스트라 알고리즘을 약간 변형하여 풀었다.

우선 다익스트라 함수는 한 지점에서 얻을 수 있는 모든 아이템의 개수를 구하는 역할이다. 우리가 알던 다익스트라에서 변형된 점은, 큐에 원소를 넣기 전에 수행하는 조건문에, 탐색범위(m)에 대한 조건을 추가하였다는 점이다. 이것으로 얻을 수 있는 dp를 돌면서, dp의 초기값(INF)이 아닐 경우 이 노드를 방문하였다는 것이다. 따라서 이 때 해당 지점에서 얻을 수 있는 아이템의 개수를 리턴할 값(possibleItemCount)에 더하였고, 이것을 리턴해주었다.

이제, 특정 지점에서 얻을 수 있는 아이템 개수를 알 수 있으므로, 이것을 모든 노드에 대해 돌아서 최대값을 구하여 출력해주면 된다.