[백준] 14938번 서강그라운드 (파이썬)

난이도 : GOLD IV

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

문제해결 아이디어

• 어디에 착지했을 경우 가장 많은 아이템을 얻는 지 확인해야하기 때문에 착지 가능 지역을 for문으로 순회하면서 해당 인덱스를 시작점으로 다익스트라 함수를 실행한다.
• 다익스트라 함수 내에서는 모든 인덱스의 위치까지 최단 경로를 구하고, 수색범위(m) 이내의 지역만 아이템 획득하고 획득한 아이템의 합을 반환한다.
  • 이 부분에서 수색범위(m) 이상의 값이면 continue를 통해서 시간 단축을 할 수 있지않을까..?
• 각 인덱스별로 획득한 아이템의 합을 비교해서 가장 큰 값을 제출한다.

import sys
import heapq

def dijk(start):
    # 각 start 인덱스를 기준으로 각각의 인덱스 까지의 거리를 모두 구한후,
    # 수색범위 이하의 거리를 가진 인덱스만 선별하여 items 에서 몇개 가질수 있는지 구한다
    # return은 구할수 있는 item의 합
    distance = [INF] * (n + 1)
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)

        if distance[now] < dist:
            continue

        for next, pay in graph[now]:
            cost = pay + dist
            if cost > m:
                continue
            
            if cost < distance[next] :
                distance[next] = cost
                heapq.heappush(q, (cost, next))
    cnt = 0
    for i,v in enumerate(distance):
        if v <= m: # 거리가 수색범위 이하일때
            cnt += items[i] #아이템 카운트 증가
    return cnt
input = sys.stdin.readline

n, m, r = map(int, input().split())
items = [0] + list(map(int, input().split()))
INF = int(1e9)
graph = [[] for _ in range(n + 1)]

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

# 1번부터 n번까지 순회 하면서 아이템 합의 최대를 구함
answer = [0] * (n + 1)

for i in range(1, n+1):
    answer[i] = dijk(i)

print(max(answer))