[SWEA] 1251_하나로 (Prim 이용)

📕 문제

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV15StKqAQkCFAYD 

SW Expert Academy

 

💡 풀이법

1. 미리 가중치를 구해서 인접 행렬의 해당 위치에 가중치를 저장해주었다.

2. 시작 노드인 0의 key 값을 0으로 바꿔준다.

2. key를 돌며 가장 작은 값을 가지고 있는 노드를 찾는다. 그리고 방문 체크를 해준다.

3.  2번에서 찾은 노드의 인접 노드들을 탐색하며 방문하지 않은 노드가 있다면 가중치와 해당 노드의 key값을 비교해서 현재 가중치가 key 값보다 작다면 key 값을 바꿔준다.

4. 2,3을 노드의 개수만큼 반복한다.

 

def find_MST(s):
    key[0] = 0

    for _ in range(N-1):
        min_idx = -1
        min_val = float('inf')
        for i in range(N):
            if not visited[i] and key[i] < min_val:
                min_idx = i
                min_val = key[i]

        visited[min_idx] = 1

        for i in range(N):
            if adj_mat[min_idx][i] and not visited[i]:
                weight = adj_mat[min_idx][i]
                if weight < key[i]:
                    key[i] = weight

T = int(input())
for tc in range(1, T+1):
    N = int(input())
    node_x = list(map(int, input().split()))
    node_y = list(map(int, input().split()))
    E = float(input())

    adj_mat = [[0 for _ in range(N)] for _ in range(N)]
    for i in range(N):
        for j in range(N):
            dist = ((node_x[i]-node_x[j])**2+(node_y[i]-node_y[j])**2)*E
            adj_mat[i][j] = dist
            adj_mat[j][i] = dist


    key = [float('inf')]*N
    visited = [0]*N

    s = 0
    find_MST(s)
    print('#{} {}'.format(tc, round(sum(key))))