[ SWEA / Python ] D4. 1251 - [S/W 문제해결 응용] 4일차 - 하나로

풀이.

각 좌표 값이 입력으로 주어진다.
각각의 점마다 거리를 계산해주고, 모든 점이 최소 거리로 연결되는 값을 찾는다.

최소 비용 신장 트리의 문제이다. 최소 비용 신장 트리 알고리즘은 우선순위 큐와 유니온 파인드 알고리즘을 사용한다.

우선순위 큐 -> heap

유니온 파인드 알고리즘의 경우 부모 노드를 지정하는 배열을 선언해주고, 노드가 연결될 때 작은 부모로 통일 시키는 것이다.

이 때 가장 작은 부모를 찾기 위해서, 노드 값과 부모 값이 다른 경우 재귀로 호출한다.

이 문제에서 모든 노드가 연결될 경우 loop를 탈출해야 하므로, union 하기 전에 부모 노드를 갱신 시켜준다.

모든 노드가 루트 노드로 연결 됐을 때 탈출시켜준다.

코드

def find(value):
    if parents[value] != value:
        parents[value] = find(parents[value])
    return parents[value]
def union(A, B):
    A = find(A)
    B = find(B)

    if A < B:
        parents[B] = A
    else:
        parents[A] = B

result = []
T = int(input())
for case in range(1, T + 1):
    ans = 0

    N = int(input())

    X = list(map(int, input().split()))
    Y = list(map(int, input().split()))

    E = float(input())

    pq = []

    for i in range(N):
        for j in range(i + 1, N):
            heappush(pq, (E * ((abs(X[i] - X[j]) ** 2 + abs(Y[i] - Y[j]) ** 2) ** 0.5) ** 2, i, j))

    parents = [i for i in range(N)]


    while pq:
        cost, start, end = heappop(pq)

        if parents.count(0) == N:
            break

        if find(start) != find(end):
            ans += cost
            union(start, end)




    result.append(f"#{case} {round(ans)}")

for r in result:
    print(r)