1774. 우주신과의 교감 (Python)

1774번: 우주신과의 교감

 

풀이

 최소 신장 트리(MST)를 이용하여 해결하는 문제이다. MST를 만드는 방법에는 크루스칼 알고리즘과 프림 알고리즘이 있다. 자세한건 검색을 통해 알아보자. 간단하게 설명하면 다음과 같다.

* 크루스칼 알고리즘 : 모든 간선을 길이에 따라 오름차순으로 정렬하여 크기가 작은 간선부터 사용한다. 간선을 사용했을 때 사이클이 발생하지 않는지를 체크한다.

* 프림 알고리즘 : 시작 정점을 기준으로 선택할 수 있는 간선 중 가장 크기가 작은 간선을 사용하여 점점 확장해 나가며 MST를 완성하는 알고리즘

 

 N은 정점의 수이고 E는 간선의 개수일 때, 크루스칼 알고리즘의 시간 복잡도는 O(E logE)이고 프림 알고리즘의 시간 복잡도는 O(N^2)이다. 이 때, 문제의 조건에 따라 알고리즘을 선택해서 사용하는 것이 좋다.

 정점에 비해 간선이 적은 그래프에서는 크루스칼 알고리즘이 적합하고, 정점에 비해 간선이 많은 그래프에서는 프림 알고리즘이 적합하다.

 

 해당 문제는 모든 좌표가 연결될 수 있는 후보에 속한다. 그렇기 때문에 간선이 많은 그래프에 속하기 때문에 프림 알고리즘을 사용하면 시간을 단축시킬 수 있다. 실제로 크루스칼 알고리즘을 사용했을 때와 프림 알고리즘을 사용했을 때 걸린 시간을 비교하면 다음과 같다.

 

[크루스칼 알고리즘]

 

[프림 알고리즘]

 

 프림 알고리즘을 사용할 때 이미 연결된 통로의 경우에는 거리를 0으로 설정해주도록 하자. 모든 좌표끼리의 거리를 구할 때 중복된 좌표끼리의 거리가 들어가지만 좌표들끼리의 거리가 0보다 작을 수 없기 때문에 거리가 0인 좌표가 선택됨이 보장된다.

 

코드

[크루스칼 알고리즘]

import sys
import heapq
import math
input = sys.stdin.readline


# 크루스칼 알고리즘을 위한 union-find
def find(n):
    if p[n] == n:
        return n
    p[n] = find(p[n])
    return p[n]


def union(a, b):
    a = find(a)
    b = find(b)

    if a == b:
        return

    p[b] = a


if __name__ == '__main__':
    N, M = map(int, input().split())
    coords = [(0, 0)] # 좌표 리스트
    p = [i for i in range(N + 1)] # 각 좌표의 부모 리스트 
    q = []

    for _ in range(N):
        a, b = map(int, input().split())
        coords.append((a, b))

    # 이미 연결된 통로는 union을 통해서 이어준다.
    for _ in range(M):
        x, y = map(int, input().split())
        union(x, y)

    # 모든 좌표에 대하여 길이를 구해준다.
    for i in range(1, N):
        for j in range(i+1, N+1):
            x = abs(coords[j][0] - coords[i][0])
            y = abs(coords[j][1] - coords[i][1])
            dist = math.sqrt(x ** 2 + y ** 2)
            heapq.heappush(q, (dist, i, j)) # 모든 간선을 최소힙에 넣어준다.

    res = 0.0
    # 이미 M개의 간선이 사용됐으므로 M이 N-1개가 될 때까지 진행한다.
    # 이미 연결된 통로는 union을 통해서 연결했기 때문에 다시 선택되는 경우는 없다.
    while M != N - 1 and len(q) != 0:
        dist, a, b = heapq.heappop(q)

        if find(a) != find(b):
            union(a, b)
            res += dist
            M += 1

    res = '{:.2f}'.format(res)
    print(res)

 

[프림 알고리즘]

import math
import sys
from collections import defaultdict
input = sys.stdin.readline
INF = float('inf')


def prim(start):
    distances[start] = 0  # 시작 정점의 거리를 0으로 설정
    res = 0.0

    # N번 반복한다.
    for _ in range(N):
        minNextNode, minDist = 0, INF

        # 현재 연결된 집합에서 아직 방문하지 않았고 거리가 가까운 좌표를 찾는다.
        for i in range(1, N+1):
            if not visited[i] and distances[i] < minDist:
                minNextNode = i
                minDist = distances[i]

        # 위에서 찾은 좌표를 방문 체크해주고 거리를 res에 더해준다.
        visited[minNextNode] = True
        res += minDist

        # 새롭게 추가된 좌표에서 연결되지 않은 좌표 중 갈 수 있는 좌표와 더 작은 거리 값을 가지고 있는 좌표를 찾아서
        # distances 배열을 갱신해준다.
        for dist, node in graph[minNextNode]:
            if not visited[node] and distances[node] > dist:
                distances[node] = dist

    return '{:.2f}'.format(res)


if __name__ == '__main__':
    N, M = map(int, input().split())
    graph = defaultdict(list)
    coords = [(0, 0)]  # 좌표 리스트
    visited = [False] * (N + 2)  # 집합에 포함된 좌표인지를 체크하기 위한 리스트
    distances = [INF for _ in range(N + 1)]  # 현재 집합에서 갈 수 있는 좌표들의 최소 거리

    for _ in range(N):
        x, y = map(int, input().split())
        coords.append((x, y))

    # 이미 연결된 통로 입력
    # 거리를 0으로 설정하여 넣어준다. 0보다 작은 거리 값은 없기 때문에 MST를 구성할 때 해당 경로가 선택됨이 보장된다.
    for _ in range(M):
        a, b = map(int, input().split())
        graph[a].append((0, b))
        graph[b].append((0, a))

    for i in range(1, N + 1):
        for j in range(i+1, N + 1):
            x = abs(coords[j][0] - coords[i][0])
            y = abs(coords[j][1] - coords[i][1])
            dist = math.sqrt(x ** 2 + y ** 2)

            graph[i].append((dist, j))
            graph[j].append((dist, i))

    # 시작 정점은 임의로 1로 설정하여 프림 알고리즘 수행
    print(prim(1))