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

문제

서준이는 아빠로부터 생일선물로 세계 지도를 받아서 매우 기뻤다. 세계 지도에서 최단 경로를 찾는 프로그램을 개발해서 아빠께 감사의 마음을 전달하려고 한다. 세계 지도는 도시를 정점으로 갖고 도시 간의 도로를 간선으로 갖는 무방향성 그래프이며(undirected graph), 도로의 길이가 간선의 가중치이다. 출발 정점 $X$에서 출발해서 $P$개의 중간 정점 중 적어도 세 개의 정점을 반드시 거친 후 도착 정점 $Z$에 도달하는 최단 거리를 구해서 우리 서준이를 도와주자.

입력

첫째 줄에 정점의 수 $N(10 \leq N \leq 100,000)$, 간선의 수 $M(10 \leq M \leq 300,000)$이 주어진다.

다음 $M$개 줄에 간선 정보 $u, v, w$가 주어지며 도시 $u$와 도시 $v$ 사이의 가중치가 정수 $w$인 양방향 도로를 나타낸다. $(1 \leq u, v \leq N, u ≠ v, 1 \leq w \leq 1,000,000)$

다음 줄에 $X Z$가 주어진다. $(1 \leq X, Z \leq N, X \neq Z)$

다음 줄에 $P$가 주어진다. $(3 \leq P \leq \min(100, N - 3))$

다음 줄에 $P$개의 서로 다른 중간 정점 $Y(1 \leq Y \leq N, X ≠ Y ≠ Z)$가 빈칸을 사이에 두고 주어진다.

출력

출발 정점 $X$에서 출발해서 $P$개의 중간 정점 중 적어도 세 개의 정점을 반드시 거친 후 도착 정점 $Z$에 도달하는 최단 거리를 출력한다. 도착 정점 Z에 도착할 수 없는 경우 $-1$을 출력한다.

풀이

브루트포스 + 다익스트라 문제이다.

$X - A - B - C - Z$ 로 가는 최단 거리를 찾아야 한다.

우선순위 큐를 활용한 다익스트라의 시간복잡도는 $O(E\log{V})$이다. 중간 정점은 최대 100개가 주어지며, 이 중 3개를 선택해야 하므로 경우의 수가 최대 ${}_{100}{\rm P}_3$ 이다. 모든 경우에서 다익스트라를 돌리면 당연히 시간초과가 난다.

$X$와 중간 정점들을 기준으로 다익스트라를 먼저 돌리고, 그 결과들을 활용해 모든 경우 중 최소 거리를 구해 주면 된다.

다익스트라를 최대 101번 돌리고, ${}_{100}{\rm P}_3 = 970,200$이다. 종합하여 시간 복잡도를 계산해 보면 $300,000\times\log{100,000} \times 101 + 970,200 = 152,470,200$ 가 나온다.

이 문제의 시간 제한은 6초이다. 바로 통과할 줄 알았다. 하지만 우리의 파이썬은 생각보다 느리다.

python3 말고 pypy3로 제출해야 한다. 다익스트라 결과를 저장할 dist를 리스트가 아닌 딕셔너리로 구현하고, min() 함수 대신 직접 if문으로 최소값을 찾는 등 최적화를 조금 진행해 주어야 겨우 풀린다.

마지막에 결과값이 INF일 땐 -1로 출력하는 것을 잊지 말자.

소스 코드

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

from heapq import heappop, heappush
from itertools import permutations
import sys

input = sys.stdin.readline
INF = float('inf')


def dijkstra(graph, start):
    n = len(graph)
    dist = [INF] * n
    dist[start] = 0
    queue = [(0, start)]
    while queue:
        path_len, v = heappop(queue)
        if path_len == dist[v]:
            for w, edge_len in graph[v]:
                if edge_len + path_len < dist[w]:
                    dist[w] = edge_len + path_len
                    heappush(queue, (edge_len + path_len, w))
    return dist


n, m = map(int, input().split())
graph = [[] for _ in range(n+1)]
for _ in range(m):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))
    graph[v].append((u, w))
x, z = map(int, input().split())
p = int(input())
mid = [int(x) for x in input().split()]

dist = {}
for node in mid:
    dist[node] = dijkstra(graph, node)
dist[x] = dijkstra(graph, x)

result = INF
for a, b, c in permutations(mid, 3):
    total = dist[x][a] + dist[a][b] + dist[b][c] + dist[c][z]
    if result > total:
        result = total
if result == INF:
    result = -1
print(result)

후기

파이썬 기준 시간 제한이 정말 빡빡한 문제였다. 이 문제를 파이썬으로 풀어본 사람이 없길래 파이썬으로 풀어 봤다.