파이썬으로 모든 점을 최소 비용으로 연결하는 프로그램 – 프림 알고리즘 풀이

문제 개요

(x, y) 형태의 좌표로 이루어진 배열 points가 주어졌다고 가정해 보겠습니다. 두 점 (xi, yi)와 (xj, yj)를 연결할 때 드는 비용은 두 점 사이의 맨해튼 거리(Manhattan Distance)이며, 공식은 다음과 같습니다.

|xi − xj| + |yi − yj|

우리가 구해야 할 값은 바로 모든 점을 서로 연결하기 위한 최소 비용입니다.

예시

입력이 다음과 같다고 가정해 봅시다.

points = [(0,0), (3,3), (2,10), (6,3), (8,0)]

이 경우 출력은 22가 됩니다. 선택된 간선들의 비용이 차례대로 6, 5, 3, 8이므로, 총 비용은 (6 + 5 + 3 + 8) = 22입니다.

접근 방식: 프림 알고리즘(Prim's Algorithm)

이 문제는 그래프 이론의 최소 신장 트리(Minimum Spanning Tree, MST) 문제와 본질적으로 같습니다. 즉, 모든 정점(점)을 연결하면서 간선 비용의 합이 최소가 되는 트리를 찾으면 됩니다. 우선순위 큐(최소 힙)를 활용한 프림 알고리즘을 사용하면 효율적으로 해결할 수 있습니다.

풀이 단계

• points_set := 0부터 (점의 개수 − 1)까지의 인덱스를 담는 새로운 집합 생성
• heap := (0, 0) 쌍으로 초기화한 최소 힙 생성
• visited_node := 방문한 노드를 추적하는 새로운 집합 생성
• total_distance := 0으로 초기화
• 힙이 비어 있지 않고, 방문한 노드 수가 전체 점의 개수보다 작은 동안 다음을 반복:
  • (distance, current_index) := 힙에서 가장 작은 원소 꺼내기
  • current_index가 아직 방문하지 않은 노드라면:
    • visited_node에 current_index 추가
    • points_set에서 current_index 제거
    • total_distance에 distance 더하기
    • (x0, y0) := points[current_index]
    • points_set의 각 next_index에 대해:
      • (x1, y1) := points[next_index]
      • (|x0 − x1| + |y0 − y1|, next_index)를 힙에 삽입

반복이 끝나면 total_distance를 반환합니다.

구현 예제

아래 파이썬 코드를 통해 더 자세히 이해해 보겠습니다.

import heapq

def solve(points):
    points_set = set(range(len(points)))
    heap = [(0, 0)]
    visited_node = set()
    total_distance = 0
    while heap and len(visited_node) < len(points):
        distance, current_index = heapq.heappop(heap)
        if current_index not in visited_node:
            visited_node.add(current_index)
            points_set.discard(current_index)
            total_distance += distance
            x0, y0 = points[current_index]
            for next_index in points_set:
                x1, y1 = points[next_index]
                heapq.heappush(heap, (abs(x0 - x1) + abs(y0 - y1), next_index))
    return total_distance

points = [(0,0),(3,3),(2,10),(6,3),(8,0)]
print(solve(points))

입력

[(0,0),(3,3),(2,10),(6,3),(8,0)]

출력

22

시간 복잡도

각 점을 방문할 때마다 아직 연결되지 않은 모든 점과의 거리를 계산하여 힙에 삽입하므로, 시간 복잡도는 O(n² log n)입니다. 점의 개수가 매우 많은 경우에는 간선을 거리 순으로 미리 정렬하거나 K-D 트리 같은 자료구조를 활용해 최적화할 수 있습니다.