[Python] SW Expert Academy #1251 하나로

📚문제: #1251 하나로(D4)

전체 코드

# 1251 하나로

from math import sqrt
from heapq import heappush, heappop

# Prim's Algorithm
def min_cost(G, t):
    
    # 전체 비용
    total = 0
    
    # 우선순위 큐
    q = []
    
    # 최소 스패닝 트리를 구성하는 노드 목록
    mst = []
    
    # 시작노드는 0, 시작노드와 연결된 간선을 큐에 추가해준다. 비용이 낮은 순서로 정렬 된다.
    for neighbor in graph[0]:
        
        heappush(q, neighbor)
    
    # 시작노드 반영
    mst.append(0)
    
    # 큐가 빌 때까지 반복한다.
    while q:
        
        cost, node = heappop(q)
        
        # 노드가 mst 구성목록에 없을 때
        if node not in mst:
            
            # 간선 비용 누적
            total += cost
            
            # mst 구성요소에 추가
            mst.append(node)
            
            if len(mst) == N:
                
                break
            
            # 해당 노드의 이웃들을 살펴보자.
            for neighbor in graph[node]:
                
                n_cost, n_node = neighbor
                
                # 이미 mst에 반영되어있으면 skip, mst에 없으면 큐에 넣는다.
                if n_node not in mst:
                    
                    heappush(q, neighbor)

    # 답안 출력
    print("#{} {}".format(t, round(total)))
    
# 테스트 케이스의 수
T = int(input())

# 각 테스트 케이스
for t in range(1, T+1):
    
    # 섬의 개수 N
    N = int(input())
    
    # 섬의 X좌표
    X = list(map(int, input().split()))
    
    # 섬의 Y좌표
    Y = list(map(int, input().split()))
    
    # 환경 부담 세율 실수 E (0 <= E <= 1)
    E = float(input())
    
    # 그래프 초기화
    graph = [[] for _ in range(N)]
    
    # 간선 정보를 추가한다.
    for i in range(N):
        
        for j in range(N):
            
            # 두 노드가 다를 때만 간선 정보를 추가한다.
            if i != j:
                
                L = sqrt((X[i]-X[j])**2 + (Y[i]-Y[j])**2)
                
                cost = E * L**2
                
                # 간선 추가 (비용, 이웃노드)
                graph[i].append((cost, j))
                
    # 함수 실행
    min_cost(graph, t)