[백준] 중량제한 1939번 (python)

출처 : https://www.acmicpc.net/problem/1939

1939번: 중량제한

풀이 방법

카카오 기출문제인 등산코스 정하기 문제랑 유사하여 알고리즘을 조금 수정해서 제출했다.
https://salmon16.tistory.com/142

다른 분들은 이분 탐색으로 풀이하던데 좋은 방법인 거 같다.

문제를 해결하기 위해 다익스트라 알고리즘을 변형해서 사용했다.

원래 다익스트라 알고리즘에선 다음 노드를 갱실 할 때 (현재 노드까지의 거리 + 현재노드에서 다음 노드의 간선의 weight)을 이용해서 다음 노드의 distance를 업데이트해주는데

 

이 문제에선 현재 노드까지의 distance와 현재 노드에서 다음 노드의 간선의 weight 중 작은 값을 다음 노드의 distance를 업데이트하는 데 사용해야 한다.

 

그리고 min heap을 사용하기 위해 weight를 음수로 변형해서 절댓값이 가장 큰 수가 pop이 되도록 설정했다.

또한 한번에 방문 가능한 경우 예외를 처리하기 위해서 다음 노드가 도착지인데 distance[도착지]가 -1인 경우에도 예외 처리를 해주었다.

import heapq
N, M = map(int, input().split())
graph = [[] for i in range(N+1)]
pq = []
start_p , end_p = 0, 0
distance = [-1 for _ in range(N+1)] ## 거리 리스트 생성
for i in range(M): ## 그래프 만들기 
    a, b, c = map(int, input().split())
    graph[a].append([b, c])
    graph[b].append([a, c])
start_p, end_p = map(int, input().split()) ## 두 공장 받기
def dijkstra(start):    
    for nn, ed in graph[start]: ## 시작 점에서 가능한 node 다 추가 
        heapq.heappush(pq, [-ed, nn]) ## 거리, 노드    

    while pq:
        weight, node = heapq.heappop(pq)        
        if node == end_p:
            if distance[node] == -1: ## 한 번에 방문가능한 경우 
                distance[node] = -weight
            return
        if distance[node] > -weight: ## 이미 더 좋은 경우가 있는 경우 탐색 x 
            continue
        for next_node, edge_weight in graph[node]:
            if next_node == start_p: ## 시작 노드이면 방문 안함
                continue
            new_weight = max(weight, -edge_weight) ## 작은 값을 선택하기 위해 음수로 둘 중 더 작은 값 선택             
            if distance[next_node] < -new_weight: ## 다음 노드보다 큰 weight으로 갈 수 있을 때 
                distance[next_node] = -new_weight
                heapq.heappush(pq, [new_weight, next_node]) ## min heap을 사용하기 위해 음수로 값을 저장하면 절댓값이 가장 큰 값이 나온
    return
dijkstra(start_p)
print(distance[end_p])