SWEA 1251 하나로

최은우·2022년 4월 5일
0

각 섬까지의 해저 터널을 만드는데 모든 섬을 연결하면서 그 해저터널의 길이가 최소가 되게끔 하는 문제.

  • 전형적인 그래프의 최소 비용 문제

풀이방법

  • 그래프 최소 비용 문제를 풀 때 정보를 사용하기 위한 밑 작업으로 infos 배열을 만들어 놓습니다.
    infos의 요소는 다음과 같은 형태를 가집니다.
    [시작 노드, 끝 노드, 두 노드 사이의 환경부담금]
    모든 노드와 노드사이를 조사해야하므로 infos는 NC2의 길이를 가지게됩니다.

  • KRUSKAL 알고리즘을 사용하기 위해 infos를 환경부담금의 오름차순으로 정렬합니다. infos.sort(key=lambda x:x[2])

  • KRUSKAL
    모든 노드의 부모노드를 저장하는 배열 p를 생성합니다.
    (초기 값은 자기 자신)

p = [0] * N
    for i in range(1, N):
        p[i] = i

kruskal 함수 내의 반복문으로 infos에서 노드 번호를 받을 때마다 부모 노드가 같은지 조사합니다 -> 같으면 cycle이므로 pass
부모 노드가 다르다면 두 노드가 들어가 있는 tree를 합치게 됩니다.

코드

from itertools import combinations

def uni(x, y):
    link(find_set(x), find_set(y))


def link(x, y):
    if rank[x] > rank[y]:
        p[y] = x
    else:
        p[x] = y
        if rank[x] == rank[y]:
            rank[y]


def find_set(n):
    if n != p[n]:
        p[n] = find_set(p[n])
    return p[n]


def kruskal(g):
    result = 0
    for info in g:
        if find_set(info[0]) != find_set(info[1]):
            result += info[2]
            uni(info[0], info[1])
    return result


T = int(input())

for tc in range(1, T + 1):
    N = int(input())  # 섬의 개수
    X_coords = list(map(int, input().split()))
    Y_coords = list(map(int, input().split()))
    E = float(input())

    infos = []
    for comb in list(combinations(list(range(N)), 2)):
        dist = (X_coords[comb[0]] - X_coords[comb[1]]) ** 2 + (Y_coords[comb[0]] - Y_coords[comb[1]]) ** 2
        cost = E * dist
        infos.append([comb[0], comb[1], cost])
    
    infos.sort(key=lambda x:x[2])
    # pprint.pprint(infos)

    p = [0] * N
    for i in range(1, N):
        p[i] = i

    rank = [0] * N

    ans = kruskal(infos)
    print(f'#{tc} {round(ans)}')
profile
코딩 아자 !

0개의 댓글