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

Thumbnail: Photo by Ugne Vasyliute (Unsplash)

간만에 그래프 문제 풀어보니까 재밌었습니다 ㅎㅎ

풀이

다익스트라(Dijkstra) vs 플로이드-워샬(Floyd Warshall)

그래프에서 최단 경로를 찾는 문제이므로, 생각할 수 있는 풀이 방법은 다익스트라(Dijkstra)와 플로이드-워샬(Floyd Warshall)의 2가지가 있습니다.

하지만, 문제의 조건에서 정점의 개수 \(V\)는 최대 1000개이므로,플로이드-워샬 알고리즘으로는 1000개의 정점에 대해 모두 최단 경로를 계산해 시간 제한을 초과할 수 있습니다. 따라서, 이 문제의 경우는 다익스트라 알고리즘이 적합합니다.

다익스트라(Dijkstra)!

다익스트라 알고리즘은 그리디하게 현재 방문한 정점들에서 가장 가까운 거리에 위치하는 정점을 찾아가는 방식입니다.

다익스트라 알고리즘의 구현에 앞서, 먼저 그래프를 구현해야 합니다. 이 문제에서는 인접 행렬 방식을 사용해 입력받은 그래프를 저장했습니다. (최대 정점의 개수가 1000개로 적기 때문.)

N, V, E = map(int, input().split())
KIST, SEA_FOOD = map(int, input().split())
INF = 101 # 간선의 가중치 값보다 큰 값을 사용한다.

# 정점은 1~V까지이므로 굳이 저장하지 않고, 이 중 팀원의 집 위치는 전체 정점의 집합의 일부이므로 별도로 저장.
Homes = list(map(int, input().split()))
# 인접 행렬 초기화. 정점의 값이 1부터 시작하므로 0번은 사용하지 않음.
Edges = [[INF for _ in range(V + 1)] for _ in range(V + 1)]

for _ in range(E):
    u, v, weight = map(int, input().split())
    Edges[u][v] = weight
    Edges[v][u] = weight

이제 다익스트라 알고리즘을 구현해 봅시다. 여기서는 우선순위 큐를 사용해 시간 복잡도를 줄였다.

💡 수정)
Python3의 queue 모듈은 멀티스레드 환경에서의 큐 처리를 위한 모듈입니다.
단순히 큐의 기능이 필요한 경우 collections모듈의 deque (double-ended queue)를 사용하는게 성능 면에서 유리합니다. 또한, 우선순위 큐 기능이 필요한 경우 heapq모듈을 활용해 list를 우선순위 큐로 사용할 수 있습니다. Python 환경에서 PS를 할 계획이시라면 이 두가지 방법이 적절합니다.

이 글은 해당 사실을 알기 전에 작성되어 PriorityQueue를 사용해 구현되었습니다. 이 글의 코드를 수정하진 않고, 대신 heapq를 활용한 다익스트라 구현을 다루는 다른 글을 참고해 주세요.

from queue import PriorityQueue

def dijkstra(start, distance):
    pq = PriorityQueue()
    pq.put((0, start))
    distance[start] = 0
    while not pq.empty():
        dist, current = pq.get()
        if distance[current] < dist:
            # 현재 우선순위 큐에서 꺼낸 상태는 최신 상태가 아님. 따라서 스킵.
            continue
        for i in range(1, V + 1):
            if distance[current] + Edges[current][i] < distance[i]:
                distance[i] = distance[current] + Edges[current][i]
                pq.put((distance[i], i))

각 팀원의 집에서 최단거리를 계산해야 할까?

문제에서는 모든 사람의 최단거리의 합을 구해야 한다.

하지만, 문제 조건에서 주어지는 그래프는 무방향 그래프이다. 그렇다면, 다익스트라 알고리즘으로 구하는 최단거리는 다르게 해석할 수 있다.

• 다익스트라 알고리즘으로 구하는 최단거리는 한 정점에서 다른 모든 정점으로의 최단 거리 이다.
• 하지만, 그래프가 무방향 그래프라면 반대로도 해석할 수 있다!
  다른 모든 정점에서 한 정점으로의 최단거리 라는 뜻이다.

이를 기반으로, 모든 팀원들의 집에서 최단거리를 계산할 필요 없이 KIST와 씨알푸드의 2개 지점에서의 최단거리만 계산해도 답을 구할 수 있다.

전체 코드

from queue import PriorityQueue
input = open(0).readline
N, V, E = map(int, input().split())
KIST, SEA_FOOD = map(int, input().split())
INF = 10101

Homes = list(map(int, input().split()))
Edges = [[INF for _ in range(V + 1)] for _ in range(V + 1)]   # 인접 행렬 초기화. 정점의 값이 1부터 시작하므로 0번은 사용하지 않음.

for _ in range(E):
    u, v, weight = map(int, input().split())
    Edges[u][v] = weight
    Edges[v][u] = weight

def dijkstra(start, distance):
    pq = PriorityQueue()
    pq.put((0, start))
    distance[start] = 0
    while not pq.empty():
        dist, current = pq.get()
        if distance[current] < dist:
            # 현재 우선순위 큐에서 꺼낸 상태는 최신 상태가 아님. 따라서 스킵.
            continue
        for i in range(1, V + 1):
            if distance[current] + Edges[current][i] < distance[i]:
                distance[i] = distance[current] + Edges[current][i]
                pq.put((distance[i], i))

res = 0
d1 = [INF for _ in range(V + 1)]
d2 = [INF for _ in range(V + 1)]
# 무방향 그래프이므로, 다익스트라를 통한 최단거리 탐색 결과는 한 점에서 모든 점까지의 거리와도 같다.
dijkstra(KIST, d1)
dijkstra(SEA_FOOD, d2)
for home in Homes:
    res += d1[home] if d1[home] != INF else -1
    res += d2[home] if d2[home] != INF else -1
print(res)