알고리즘/BOJ

[BOJ/python]1162번 도로포장

frog-in-well 2022. 7. 6. 22:51

백준 1162번 도로포장 파이썬, 다익스트라

문제 링크 - https://www.acmicpc.net/problem/1162

 

1162번: 도로포장

첫 줄에는 도시의 수 N(1 ≤ N ≤ 10,000)과 도로의 수 M(1 ≤ M ≤ 50,000)과 포장할 도로의 수 K(1 ≤ K ≤ 20)가 공백으로 구분되어 주어진다. M개의 줄에 대해 도로가 연결하는 두 도시와 도로를 통과하

www.acmicpc.net

주어진 도로에서 최대 K개의 도로의 가중치를 0으로 만들 수 있다. 조합을 사용할까 고민했다. 다익스트라를 이용하여 dist배열을 2차원으로 만들고 도로의 가중치를 0으로 만든 횟수를 저장했다.

1. 풀이

  • heap에 현재까지 누적 가중치, 0으로 만든 도로의 수, 현재 노드를 넣어준다.
  • 다음 노드를 방문할 때 heap에서 pop된 정보를 이용하여 최단 거리를 구한다.
  • 다음 노드를 방문할 때 0으로 만든 도로의 수를 기준으로 방문 표시를 한다.

2. 코드

from heapq import heappop, heappush

INF = float('inf')
N,M,K = map(int, input().split())
G=[ [] for _ in range(N+1)]
for _ in range(M):
    u,v,c = map(int,input().split())
    G[u].append([v,c])
    G[v].append([u,c])
    
def dijkstra():
    
    dist=[ [INF for _ in range(K+1)] for _ in range(N+1)]
    visited=[[0 for _ in range(K+1)] for _ in range(N+1)]
    heap=[]
    dist[1] = [0 for _ in range(K+1)]
    heappush(heap, [0,0,1])
    
    while heap:
        cum, deleted,node = heappop(heap) 
        if visited[node][deleted]==1:
            continue
        
        visited[node][deleted]=1
 
        
        for next, cost in G[node]:
            if visited[next][deleted]==1:
                continue
            
            if deleted<K :
                if dist[next][deleted+1] > cum:
                    dist[next][deleted+1]=cum
                    heappush(heap, [cum, deleted+1, next])
                
            if deleted==0:
                if dist[next][deleted] > cum + cost:
                    dist[next][deleted] =cum + cost
                    heappush(heap, [dist[next][deleted], deleted, next])
                
            else:
                if dist[next][deleted] > cum + cost or (dist[next][deleted] > dist[node][deleted-1]+0 ):
                    
                    if cum+ cost > dist[node][deleted-1] :
                        dist[next][deleted] = dist[node][deleted-1]
                        heappush(heap, [dist[next][deleted], deleted, next ])
                    else :
                        dist[next][deleted] = cum + cost
                        heappush(heap, [dist[next][deleted], deleted, next])
                    
    return dist

result = dijkstra()
print(min(result[N]))