[BOJ][Python]14938 풀이

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

14938번: 서강그라운드

import sys; ssr = sys.stdin.readline
import heapq
INF = 16

n, m, r = map(int, ssr().split())
items = list(map(int, ssr().split()))
routes = [[] for _ in range(n+1)]
ans = 0
for _ in range(r):
    a, b, l = map(int, ssr().split())
    routes[a].append((b, l))
    routes[b].append((a, l))
for start in range(1, n+1):
    q = [(0, start)]
    min_dist = [INF for _ in range(n+1)]
    min_dist[start] = 0
    while q:
        cur_dist, cur_node = heapq.heappop(q)
        if cur_dist > min_dist[cur_node]:
            continue
        for next_node, next_dist in routes[cur_node]:
            dist = cur_dist + next_dist
            if dist < min_dist[next_node]:
                min_dist[next_node] = dist
                heapq.heappush(q, (dist, next_node))
    item_cnt = 0
    min_dist = min_dist[1:]
    for i in range(len(min_dist)):
        if min_dist[i] <= m:
            item_cnt += items[i]
    if item_cnt > ans:
        ans = item_cnt
print(ans)

문제의 조건이 착지점 기준 이동 범위 내 아이템을 수집하는 것이기 때문에 출발점에서의 모든 노드까지의 거리를 안다면 쉽게 구할 수 있는 문제라는 생각이 드실 겁니다. '출발점에서 모든 노드까지의 거리' 라는 말을 들으면 감이 오시는 분들도 있으실겁니다. 데이크스트라를 사용하라는 얘기죠.

 

그래서 데이크스트라를 이용해서 풀이를 작성해봤습니다. 다만 위 문제는 출발점을 명시하지 않고 수집할 수 있는 아이템의 최대 갯수를 물어봤기 때문에 모든 노드를 출발점으로 해서 각 노드까지의 거리를 구한 다음 아이템의 갯수를 비교해야 합니다. 다시 말하면 모든 노드에서 다른 모든 노드 간의 거리를 구하면 된다는 얘기죠. 플로이드-워셜을 사용해도 되겠구나 하는 생각이 드는 문장이죠.

 

import sys; ssr = sys.stdin.readline
INF = 16

n, m, r = map(int, ssr().split())
items = list(map(int, ssr().split()))
routes = [[] for _ in range(n+1)]
min_dist = [[INF for _ in range(n+1)] for _ in range(n+1)]
for _ in range(r):
    a, b, l = map(int, ssr().split())
    routes[a].append((b, l))
    routes[b].append((a, l))
    min_dist[a][b] = min(min_dist[a][b], l)
    min_dist[b][a] = min(min_dist[b][a], l)
for i in range(n+1):
    min_dist[i][i] = 0
for mid in range(n+1):
    for start in range(n+1):
        for end in range(n+1):
            min_dist[start][end] = min(min_dist[start][end], min_dist[start][mid]+min_dist[mid][end])
ans = 0
for i in range(n+1):
    tmp = 0
    for j in range(n+1):
        if min_dist[i][j] <= m:
            tmp += items[j-1]
    if ans < tmp:
        ans = tmp
print(ans)

 

플로이드-워셜로도 한 번 작성해봤습니다. 그래서 위의 문제는 데이크스트라나 플로이드-워셜 둘 다 사용가능합니다. 복잡도도 $O(N^3)$으로 같을테니 런타임도 비슷할 것 같네요.