본문 바로가기

Code Review/BaekJoon

[ 백준 11265번 끝나지 않는 파티 ] Python 코드 (다익스트라, 벨만포드)

1. 문제

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

 

11265번: 끝나지 않는 파티

입력의 첫 번째 줄에는 파티장의 크기 N(5 ≤ N ≤ 500)과 서비스를 요청한 손님의 수 M(1 ≤ M ≤ 10,000) 이 주어진다. 각각의 파티장은 1번부터 N번까지 번호가 붙여져 있다. 다음에는 N개의 줄에 걸

www.acmicpc.net


2. 풀이

2.1. Bellman-Ford

입력이 2차원 배열인데다가, 출발지도 여러개여서 Bellman-Ford 알고리즘을 쓰고 싶은 충동에 휩싸인다.

실제로 Bellman-Ford를 써서 풀 수 있기는 하다! 그런데 파이썬의 비애로 자잘한 최적화를 해줘야 한다.

 

관련 내용은 이 블로그를 참조했다. 아주 자세하고 친절하게 써주셨다.

  1. 무조건 min으로 최솟값 업데이트 하지 말고, if 조건문 걸어주기
  2. for i in range에서 i 안 쓸 경우, for _ in range 로 바꾸기

 

참고로 벨만포드의 시간복잡도는 O(VE)인데, 이 문제는  E = V(V-1)인 완전 무방향 그래프이므로 O(V^3)이 된다.


2.2. Dijkstra

Dijkstra의 시간복잡도가 O(V^2)인데, 이를 각 노드에 대해서 V번 수행하면,

O(V^3)가 되므로 벨만포드 알고리즘과 시간복잡도가 동일해진다.

 

그런데 생각해보면, 반드시 모든 노드에 대해서 최단경로를 계산할 필요는 없다.

입력에서 주어진 출발지들에 대해서만 최단경로를 계산하면 된다.

즉 노드가 100개라도 주어진 입력에서의 출발지 종류가 20개라면 다익스트라를 20번만 수행하면된다.

출발지 목적지
1 54
2 15
2 3
2 76
1 11

예를 들어, 위와 같이 입력이 주어졌다면 다익스트라를 2번만 돌리면 된다.

첫번째 입력에서, 출발지 1에 대해서 다익스트라를 연산하고 그 결과를 저장해둔다.

두번째 입력에서, 출발지 2에 대해서 다익스트라를 연산하고 그 결과를 저장해둔다.

세번째 입력부터는 저장된 값을 읽어오기만 하면된다.

 

다익스트라에 힙 구조까지 적용하면, 기존의 시간복잡도 O(V^2)가 O(E*logV)로 줄어든다.


3. 코드

3.1. Bellman-Ford

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(N)]

for k in range(N):
    for n1 in range(N):
        for n2 in range(N):
            # arr[n1][n2] = min(arr[n1][n2], arr[n1][k] + arr[k][n2])
            if arr[n1][n2] > arr[n1][k] + arr[k][n2]:   # 시간개선: 위에 코드로 무작정 min연산 하지 말고, 조건문을 걸어준다.
                arr[n1][n2] = arr[n1][k] + arr[k][n2]

result = []
for _ in range(M):  # 시간개선: for i in range 대신 for _ in range를 쓴다.
    s, e, t = map(int, input().split())
    if arr[s-1][e-1] <= t:
        print('Enjoy other party')
    else:
        print('Stay here')

3.2. Dijkstra

record = [[] for _ in range(N + 1)]리스트를 만들었다.

그리고 어떤 입력이 s를 출발점으로 한다면, 먼저 record[s]를 찾고 비었을 경우에만 다익스트라를 수행하도록 했다.

다익스트라를 수행했다면 종료하기 전에, 결과값을 record[s]에 저장했다.

그러면 다음 번에 s를 출발지로 하는 입력이 주어졌을 때, 이번에는 record[s]의 값을 읽기만 하면 된다.

import sys
import heapq
input = sys.stdin.readline

def dijkstra2_2(s, e):
    INF = 1000000001
    distance = [INF] * (N + 1)              # 출발지로부터 모든 노드까지의 최단경로를 저장할 리스트
    distance[s] = 0                         # 출발지 자기 자신까지는 비용이 0

    heap = []
    heapq.heappush(heap, (0, s))            # 출발지 노드를 초깃값으로 입력 (비용, 노드)

    while heap:
        dist, now = heapq.heappop(heap)     # 주변 노드 중 가장 비용이 짧은 노드를 pop
        if distance[now] < dist:            # 현재 노드가 처리된적 있다면 무시
            continue
        for next in range(1, N + 1):        # 현재 노드의 주변 노드를 탐색
            cost = dist + arr[now][next]        # 현재 노드를 거치는 값을 임시로 계산
            if cost < distance[next]:           # 현재 노드를 거치는 것이 더 빠를 경우
                distance[next] = cost               # 해당 노드의 최단경로 갱신
                heapq.heappush(heap, (cost, next))  # 힙에 노드정보 입력 (비용, 노드)

    record[s] = distance	# 다익스트라 끝내기 전에 결과값 record[s]에 저장!!
    return distance[e]

N, M = map(int, input().split())
arr = [[]] + [[0] + list(map(int, input().split())) for _ in range(N)]

record = [[] for _ in range(N + 1)]
for _ in range(M):
    s, e, t = map(int, input().split())
    if record[s]:                       # record[s]가 비어있지 않다면 바로 값을 읽어와서 결과를 출력한다.
        if record[s][e] <= t:               
            print('Enjoy other party')
        else:
            print('Stay here')
    else:                               # record[s]가 비어있다면 다익스트라 연산을 수행한다.
        if dijkstra2_2(s, e) <= t:
            print('Enjoy other party')
        else:
            print('Stay here')