Baekjoon 11265. 끝나지 않는 파티 / Python

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

문제설명

문제는 놀이동산 "민호월드"의 여러 파티장에서 일방통행 도로를 통해 다른 파티장으로 이동할 때, 특정 시간 내에 도착할 수 있는지를 판단하는 것입니다.

1. 각 파티장은 다른 모든 파티장과 직접적으로 연결된 도로가 있습니다.
2. 각 도로는 일방통행이며, 각 도로를 통해 이동하는 시간이 주어집니다.
3. M개의 쿼리에서 출발 파티장, 도착 파티장, 그리고 시간 제한이 주어지며, 주어진 시간 내에 도착할 수 있는지를 판단해야 합니다.

코드 설명

사용된 알고리즘 : 다익스트라

결과 리스트 res를 만들어 중복을 관리한다.

이미 계산된 결과를 res 리스트에서 바로 사용하기 위함이다.

처음 계산하는 값이라면,

다익스트라 함수를 통해 res에 최단 경로를 계산한다.

우선순위 큐를 사용해서 최소 비용과 위치를 저장한다.

import sys;input = sys.stdin.readline
from heapq import heappop, heappush

# 파티장 수 N 쿼리 수 M
n, m = map(int, input().split())

# 파티장 간의 이동 시간
road = [list(map(int, input().split())) for _ in range(n)]

inf = 1e10

# 각 파티장에서 다른 파티장까지의 최단 경로를 저장할 리스트
res = [[] for _ in range(n)]

# 다익스트라 알고리즘 함수
def sol(s, e):
    # 만약 출발지와 도착지가 같다면 이동 시간은 0
    if s == e:
        return 0
    
    # 최단 거리를 저장할 리스트
    dist = [inf] * n
    dist[s] = 0  # 시작점

    # 우선순위 큐
    q = []
    heappush(q, (0, s))  # 시작점, 거리 0
    
    while q:
        # 현재 가장 작은 비용과 그에 해당하는 파티장
        cost, now = heappop(q)
        
        # 만약 저장된 거리가 현재 비용보다 작으면 건너뜀
        if dist[now] < cost:
            continue
        
        # 현재 파티장에서 다른 모든 파티장으로의 이동 고려
        for next in range(n):
            tmp = cost + road[now][next]
            
            # 만약 새로운 경로가 기존 경로보다 짧으면 갱신
            if tmp < dist[next]:
                dist[next] = tmp
                heappush(q, (tmp, next))  # 새로운 경로를 우선순위 큐에 추가
    
    # 현재 파티장에서 다른 모든 파티장으로의 최단 경로 저장
    res[s] = dist
    
    # 도착지까지의 최단 경로 반환
    return dist[e]


for _ in range(m):
    # 출발 파티장(s), 도착 파티장(e), 제한 시간(t)
    s, e, t = map(int, input().split())
    
    # 만약 출발지에서의 최단 경로가 이미 계산되어 있다면
    if res[s - 1]:
        # 도착지까지의 최단 경로가 제한 시간 내에 있는지 확인
        if res[s - 1][e - 1] <= t:
            print('Enjoy other party')
        else:
            print('Stay here')
    # 처음 계산하는 경우
    else:
        # 다익스트라 알고리즘을 호출하여 최단 경로 계산
        if sol(s - 1, e - 1) <= t:
            print('Enjoy other party')
        else:
            print('Stay here')