Sign in Subscribe
PS

[PS] BOJ 14938 / 서강그라운드

Minjun Kim

• 2 min read
[PS] BOJ 14938 / 서강그라운드
Thumbnail: Photo by Yurii Khomitskyi (Unsplash)
문제 링크: https://www.acmicpc.net/problem/14938

다익스트라 문제입니다.

풀이

다익스트라

​오랜만에 푸는 다익스트라 문제인 만큼, 다익스트라 알고리즘의 개념을 다시 정리해 봅시다.

다익스트라 알고리즘은 정점을 기준으로 최단 경로를 찾는 알고리즘이며, 한 개의 정점에서 다른 모든 정점으로의 최단 경로를 계산합니다. 아래는 우선순위 큐 기반의 다익스트라 구현입니다.

from heapq import heappop, heappush

graph = [[] for _ in range(N + 1)]
... # 그래프 입력

def dijkstra(start):
    distance = [INF] * (N + 1)
    distance[start] = 0
    pq = [(0, start)]
    
    while pq:
        dist, vertex = heappop(pq)
        if dist < distance[vertex]:
            continue

        for v, c in graph[vertex]:
            if dist + c < distance[v]:
                distance[v] = dist + c
                heappush(pq, (dist + c, v))
    
    return distance

우선순위 큐 기반의 다익스트라 구현

문제 풀이

이 문제는 $1 \cdots N$의 정점 중 어디에서 출발해야 거리 $M$ 이내로 도달할 수 있는 모든 정점에 담긴 아이템의 수의 합이 최대가 되는지 구하는 문제입니다.

따라서, 다음 순서로 풀면 됩니다.

  1. $1 \cdots N$까지의 정점들에서 출발하는 최단 거리를 다익스트라를 통해 구한다.
    1. 다익스트라를 통해 $i$번 정점에서의 최단 거리를 구했다면, 정점 $i$에서 거리 $M$ 이내로 도달할 수 있는 정점을 찾는다.
    2. a에서 찾은 정점에 대해, 해당 정점에 담긴 아이템의 수의 합을 구한다.
  2. 1에서 구한 아이템의 수의 합의 최댓값을 구한다.

전체 코드

from heapq import heappush, heappop
input = open(0).readline

INF = 20
N, M, R = map(int, input().split())
items = list(map(int, input().split()))
graph = [[] for _ in range(N + 1)]
pq = []
for _ in range(R):
    a, b, l = map(int, input().split())
    graph[a].append((b, l))
    graph[b].append((a, l))

def dijkstra(start):
    distance = [INF] * (N + 1)
    distance[start] = 0
    pq.clear()
    heappush(pq, (0, start))
    while pq:
        dist, vertex = heappop(pq)
        if dist < distance[vertex]:
            continue

        for v, c in graph[vertex]:
            if dist + c < distance[v]:
                distance[v] = dist + c
                heappush(pq, (dist + c, v))
    
    cnt = 0
    for v in range(1, N + 1):
        if distance[v] <= M:
            cnt += items[v - 1]
    return cnt

max_items = 0
for i in range(1, N + 1):
    max_items = max(max_items, dijkstra(i))
print(max_items)

solution.py