백준 1238: 파티(python)

문제

---

풀이

 

각 집에 대해, X에서 집으로 가는 경로와 집에서 X로 가는 경로를 더한 값이 최대인 집을 찾아야 한다. X->집과 집->X를 모두 구해야 하므로 모든 지점에서 다른 모든 지점까지의 최단 거리를 구하기 위해 플로이드-워셜 알고리즘을 사용하였으나, 플로이드-워셜 알고리즘의 시간 복잡도는 노드 개수 V에 대해 \(O(V^3)\)이므로 시간 초과 오류가 발생한다.

플로이드-워셜 코드 참고

import sys
input = sys.stdin.readline

N, M, X = map(int, input().split()) # N=학생(노드), M=도로(간선), X=목적지(노드)

dist = [[sys.maxsize for _ in range(N + 1)] for _ in range(N + 1)]
for i in range(N + 1):
    dist[i][i] = 0   # init 대각선

for i in range(M):
    s, e, t = map(int, input().split())
    dist[s][e] = t

for k in range(1, N + 1):
    for s in range(1, N + 1):
        for e in range(1, N + 1):
            dist[s][e] = min(dist[s][e], dist[s][k] + dist[k][e])

res = [0 for _ in range(N + 1)]

for i in range(1, N + 1):
    res[i] = dist[X][i] + dist[i][X]

print(max(res))

따라서 \(O(ElogV)\)인 다익스트라 알고리즘을 사용하였으며, 다익스트라 알고리즘은 한 지점으로부터 다른 모든 지점까지의 최단 거리를 탐색하는 알고리즘이므로 집->X와 X->집 양방향의 경로를 모두 저장하기 위해 graph배열과 reverse_graph 배열 2개를 사용한다.

N, M, X = map(int, input().split()) # N=학생(노드), M=도로(간선), X=목적지(노드)
graph = [[] for _ in range(N + 1)]          # X -> 집
reverse_graph = [[] for _ in range(N + 1)]  # 역순

for _ in range(M):
    s, e, t = map(int, input().split())
    graph[s].append((e, t))			# 순방향 저장
    reverse_graph[e].append((s, t))	# 역방향 저장

이로써 graph[i]는 X에서 i까지의 최단 거리가, reverse_graph[i]에는 i에서 X까지의 최단 거리가 저장될 것이다.

 

모든 노드(집)에 대해 집->X 최단 경로를 탐색하기 위해 다익스트라를 여러 번 호출해야 하므로 다익스트라 알고리즘은 함수로 구현해 주었으며,

def dijkstra(start, graph):
    dist = [sys.maxsize for _ in range(N + 1)]
    visited = [False for _ in range(N + 1)]
    pq = PriorityQueue()

    dist[start] = 0
    pq.put((0, start))  # pq: (시간, 노드)

    while pq.qsize() > 0:
        time, stt = pq.get()

        if visited[stt]:
            continue
        visited[stt] = True

        for end, nxt_time in graph[stt]:
            if dist[end] > dist[stt] + nxt_time:
                dist[end] = dist[stt] + nxt_time
                pq.put((dist[end], end))

    return dist

dijkstra() 함수는 시작 지점과 탐색을 진행할 배열을 인자로 받아 start로부터 각 지점까지의 거리 배열 dist를 반환한다.

in_dist = dijkstra(X, graph)    # X -> 집
out_dist = dijkstra(X, reverse_graph)   # 집 -> X

X에서 집들까지의 거리를 저장할 in_dist 배열과 집들에서 X까지의 거리를 저장할 out_dist 배열을 선언하였으며,

res = [in_dist[i] + out_dist[i] for i in range(1, N + 1)]

진입 거리와 진출 거리의 총합을 res 배열에 저장하면 res 중 값이 가장 큰 것이 총 이동 시간이 가장 긴 학생이 움직이는 시간이 된다.

 

전체 코드 

import sys
from queue import PriorityQueue
input = sys.stdin.readline

N, M, X = map(int, input().split()
graph = [[] for _ in range(N + 1)]    
reverse_graph = [[] for _ in range(N + 1)]  

for _ in range(M):
    s, e, t = map(int, input().split())
    graph[s].append((e, t))
    reverse_graph[e].append((s, t))

def dijkstra(start, graph):
    dist = [sys.maxsize for _ in range(N + 1)]
    visited = [False for _ in range(N + 1)]
    pq = PriorityQueue()

    dist[start] = 0
    pq.put((0, start))

    while pq.qsize() > 0:
        time, stt = pq.get()

        if visited[stt]:
            continue
        visited[stt] = True

        for end, nxt_time in graph[stt]:
            if dist[end] > dist[stt] + nxt_time:
                dist[end] = dist[stt] + nxt_time
                pq.put((dist[end], end))

    return dist

in_dist = dijkstra(X, graph) 
out_dist = dijkstra(X, reverse_graph)  

res = [in_dist[i] + out_dist[i] for i in range(1, N + 1)]
print(max(res))