[SWEA/Python] 1251. 하나로

📌 문제

https://swexpertacademy.com/main/main.do

SW Expert Academy

 

📌 문제 접근 방법

1. Prim 알고리즘을 사용하였다.
2. 임의의 섬을 하나 선택하고 방문체크한다.
3. 아직 방문하지 않은 섬을 순회하면서 현재 섬과의 거리와 이미 dist에 저장된 거리를 비교해서 더 작은 값을 dist에 저장한다.
4. 방문하지 않은 섬을 순회하면서 dist가 최소인 섬을 고르고 방문체크한다.
5. 4번의 섬이 2번의 섬과 연결된 것이므로 4번의 섬에 저장된 dist를 E를 곱해서 pay에 더해준다.
6. 모든 섬을 방문할 때까지 위의 과정 반복

 

📌 코드

import sys
sys.stdin = open('input.txt')

def link(start):
    global pay

    dist[start] = 0

    for _ in range(N):
        min_idx = -1
        min_val = float('inf')

        for i in range(N):
            if not visited[i] and dist[i] < min_val:
                min_idx = i
                min_val = dist[i]
        
        visited[min_idx] = 1
        pay += E * min_val

        for j in range(N):
            if not visited[j]:
                dist[j] = min(dist[j], (x[min_idx] - x[j])**2 + (y[min_idx] - y[j])**2)
        

T = int(input())

for t in range(1, T+1):
    N = int(input())
    x = list(map(int, input().split()))
    y = list(map(int, input().split()))
    E = float(input())

    visited = [0 for _ in range(N)]
    dist = [float('inf') for _ in range(N)]
    pay = 0

    link(0)

    print('#{} {}'.format(t, round(pay)))