17182. 우주 탐사선 (Python)

17182번: 우주 탐사선

 

풀이

 먼저 모든 지점에서 다른 모든 지점까지의 최단 거리를 구해야한다. 이를 위해서 플로이드 워셜 알고리즘을 사용하면 된다. 플로이드 워셜 알고리즘을 사용하여 최단 거리를 구해놓고 이후에는 순열을 통해서 모든 행성을 방문하는 경우의 수에 대해서 거리를 계산하고 그중 최단 거리를 출력하면 된다.

 

 순열을 구할 때는 시작 행성 위치는 제외하여 구해주면 된다. 문제의 2번째 예제를 예로 들면 1번 행성에서 출발하므로 총 [(0, 2, 3), (0, 3, 2), (2, 0, 3), (2, 3, 0), (3, 0, 2), (3, 2, 0)]에 대해서 거리를 구해주면 된다.

 

코드

from itertools import permutations

INF = float('INF')

# 플로이드 워셜 알고리즘 
def floyd():
    for k in range(N):
        for i in range(N):
            for j in range(N):
                graph[i][j] = min(graph[i][j], dist[i][k] + dist[k][j])


def solve():
    perList = list(permutations([x for x in range(N) if x != K],N-1)) # 출발 위치를 제외한 순열 리스트
    res = INF 

    for per in perList:
        distSum = 0 # 현재 경로에 대한 거리의 합 
        now = K # 출발 위치
        
        for p in per: # now에서 p로 가는 거리를 더해나간다.
            distSum += graph[now][p]
            now = p
        
        res = min(res, distSum) # res 갱신
    
    print(res)


if __name__ == '__main__':
    N, K = map(int, input().split()) # N: 행성의 개수, K: 발사되는 행성의 위치
    dist = [] # 입력으로 주어지는 2차원 행렬
    graph = [[INF for _ in range(N)] for _ in range(N)] # graph[a][b]: a에서 b로 가는데 걸리는 최단거리
    
    for _ in range(N):
        dist.append(list(map(int, input().split())))

    floyd() # 플로이드 워셜 알고리즘 수행
    solve() # 순열을 통해서 모든 경우의 수 탐색