백준 22870 산책(large) (파이썬)

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

22870번: 산책 (large)

풀이

이 문제는 다익스트라를 응용해서 푸는 문제이다. 문제에서 요구하는 풀이방법은 "시작 정점 ~ 도착 정점"과 "도착 정점 ~ 시작 정점"의 최소 경로가 서로 겹치면 안된다 이다.

이를 차근차근 풀어서 설명하면 다음과 같다.

1. 다익스트라 알고리즘을 통해 출발점 S에서 도착점 E까지 가는 최소 경로(dist_s)와 도착점E부터 출발점 S까지 가는 최소경로(dist_e)를 구한다.
2. 이후 경로가 겹치지 않게 하기위해 정점을 지워야 한다. 그리고 cur_node를 임의로 지정하고 cur_node와 연결되어 있는 next_node, 거리인 next_dist를 구하고 dist_s[cur_node] + next_dist + dist_e[next_node]를 구한다.
3. 만약 3에서 구한 값이 dist_s[E]와 같다면 이는 출발 ~ 도착 경로와 도착 ~ 출발 경로와 같다고 판단하고 중간 노드인 next_node들 중에서 가장 작은 노드번호를 구한다. (문제에서 사전순으로 먼저 오는 노드를 먼저 선택한다고 하였기 때문)
4. 가장 작은 노드를 방문했다고 처리한다.
5. cur_node가 도착노드를 가르킬때까지 3~4을 반복한다.
6. 도착 지점(E)에서 출발 지점(S)까지 다익스트라 알고리즘을 한번더 돌려 도착점 ~ 시작점까지의 최소 경로를 다시 구한다. (1번과의 다른점은 시작 정점 ~ 도착 정점까지 경유한 노드들을 지운뒤 구한 최소경로라는 점이다.)

코드

import heapq, sys
input = lambda: sys.stdin.readline().rstrip()

visited = [False] * 200001


def dijkstra(start):
    dp = [20000001] * (N + 1)
    pq = [(0, start)]
    dp[start] = 0

    while pq:
        cur_dist, cur_node = heapq.heappop(pq)

        for next_node, next_dist in graph[cur_node]:
            if visited[next_node]:
                continue
            if dp[next_node] > cur_dist + next_dist:
                dp[next_node] = cur_dist + next_dist
                heapq.heappush(pq, (dp[next_node], next_node))

    return dp


def erase_node(start, end):
    cur_node = start
    pre_node = start

    while cur_node != end:
        min_node = 200001
        for next_node, next_dist in graph[cur_node]:
            if next_node == pre_node:
                continue

            if dist_s[cur_node] + next_dist + dist_e[next_node] == dist_s[end]:
                min_node = min(min_node, next_node)

        visited[min_node] = True
        pre_node = cur_node
        cur_node = min_node

N, M = map(int, input().split(" "))
graph = [[] for _ in range(N+1)]
for _ in range(M):
    s, e, dist = map(int, input().split(" "))
    graph[s].append((e, dist))
    graph[e].append((s, dist))
S, E = map(int, input().split(" "))
dist_s = dijkstra(S)
dist_e = dijkstra(E)
answer = dist_s[E]
erase_node(S,E)
dist_e = dijkstra(E)
answer += dist_e[S]
print(answer)

Python3로 하면 시간초과가 난다. 그렇기에 PyPy3로 제출해야 한다.