[Python] 최소 신장 트리(MST) - 크루스칼, 프림 알고리즘

최소 신장 트리란?

가장 적은 가중치의 간선을 사용, 사이클이 생기지 않은 채 모든 정점을 연결

E = V - 1 (모든 간선에 한번씩 연결되므로)

https://velog.io/@jjhjjh1159/%ED%81%AC%EB%A3%A8%EC%8A%A4%EC%B9%BC-VS-%ED%94%84%EB%A6%BC-%ED%8C%8C%ED%97%A4%EC%B9%98%EA%B8%B0-pzljatpq

크루스칼(Kruskal) 알고리즘

Greedy method를 이용. 모든 정점을 포함하면서, 간선의 가중치 합이 최소가 되는 트리를 구함

• 최소 비용의 간선으로 구성, 사이클을 포함하지 않음. 이 두 조건에 근거하여 사이클을 이루지 않는 최소 비용 간선을 선택

# 유니온 파인드(Union-Find) 클래스
class UnionFind:
    def __init__(self, n):
        # 1번부터 n번 노드까지 부모를 자기 자신으로 초기화
        self.parent = list(range(n + 1))

    def find(self, x):
        # 경로 압축(Path Compression)을 적용하여 부모 노드를 찾음
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        # 두 노드의 루트 노드를 찾아서 같은 집합에 속하는지 확인
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a == root_b:
            return False  # 이미 같은 집합이면 union하지 않음 (사이클 발생)
        # 보통은 작은 번호의 루트를 부모로 연결하는 등의 최적화를 할 수 있음
        self.parent[root_b] = root_a
        return True

# 크루스칼 알고리즘 함수
def kruskal(n, edges):
    # 간선을 가중치 기준으로 오름차순 정렬
    edges.sort(key=lambda x: x[2])
    uf = UnionFind(n)
    mst_cost = 0
    mst_edges = []  # MST에 포함된 간선 리스트

    # 간선을 하나씩 확인하면서 사이클이 발생하지 않으면 MST에 추가
    for u, v, cost in edges:
        if uf.union(u, v):
            mst_cost += cost
            mst_edges.append((u, v, cost))

    return mst_cost, mst_edges

# 예제 사용
if __name__ == '__main__':
    # 예제: 6개의 노드와 여러 간선 정보 (u, v, cost)
    n = 6
    edges = [
        (1, 2, 3),
        (1, 3, 2),
        (2, 3, 1),
        (2, 4, 4),
        (3, 4, 5),
        (3, 5, 6),
        (4, 5, 7),
        (4, 6, 8),
        (5, 6, 9)
    ]

    total_cost, mst_edges = kruskal(n, edges)
    print("MST 총 비용:", total_cost)
    print("MST 간선들:")
    for u, v, cost in mst_edges:
        print(f"{u} - {v} : {cost}")

흐름

1. 간선들을 가중치 기준 오름차순으로 정렬
2. 정렬된 간선 리스트를 순회
   1. 두 정점을 Union-Find를 통해 사이클이 형성 되어있는지 확인
   2. 없다면, 해당 간선을 MST에 포함시키고 두 집합을 합침
3. 모든 정점이 연결될 때 까지 반복

특징

• 시간 복잡도는 간선을 정렬하는 시간에 좌우됨.
• 간선의 개수가 E인 그래프의 정렬을 퀵 정렬 등 효율적인 알고리즘으로 으로 정렬한다고 했을 때의
  시간 복잡도는 O(E log E)이 된다

의사코드

크루스칼(G):
    입력: G = (V, E)  
           V는 정점 집합, E는 간선 집합 (각 간선은 (u, v, weight) 형태)
    출력: MST (최소 신장 트리)

    MST ← {}                           // 결과로 반환할 최소 신장 트리를 저장할 집합

    // 1. 간선들을 가중치(weight) 오름차순으로 정렬
    E ← 정렬된 E (weight 기준)

    // 2. 모든 정점을 독립적인 집합으로 초기화 (Union-Find 자료구조 사용)
    for each 정점 v in V:
        MakeSet(v)

    // 3. 정렬된 간선들을 순회하며 MST 구성
    for each 간선 (u, v, weight) in E:
        if Find(u) ≠ Find(v):         // u와 v가 서로 다른 집합에 있다면 (사이클 발생 X)
            MST에 (u, v, weight)를 추가
            Union(u, v)               // u와 v를 같은 집합으로 합친다.

            if MST에 추가된 간선 수 == |V| - 1:
                break                 // 모든 정점이 연결되었으므로 종료

    return MST

Union-Find 자료구조

서로소 부분집합들로 나뉜 원소들을 처리하기 위한 자료구조.

https://doing7.tistory.com/82

Kruskal algorithm에서는 Cycle 생성 여부 판단용으로 사용

서로소 집합(Disjoint Set)?

서로 공통된 원소를 가지지 않은 두 개 이상의 집합

주요 연산

• MakeSet(x) : 모든 원소를 각각의 집합으로 만듬
  • 각 원소는 자기 자신을 대표로 만듬
  • node[x] = x

  def make_set(x):
        node[x] = x

• Find(x) : 특정 노드의 원소(대표)를 반환.
  
  • 기본적으로는 자기 자신이 대표이기 때문에 그대로 반환, 만일 Union을 통해 합쳐진 상태이면 대표의 값을 반환

  def find(node) : 
          if graph[node] != node:
                  graph[node] = find(graph[node])
          return graph[node]

• Union(x, y) : 두 집합을 하나의 집합으로 합침
  • find(x) find(y) 함수 호출, 각각의 대표를 확인
  • 대표가 다르다면, 서로 다른 집합이므로 합칠 수 있음
  • 둘중 한 노드의 원소를 대표로 변경(보통 작은값을 대표로)

  def union(x, y) :
        x_parent = find(x) # x 노드의 대표 확인
        y_parent = find(y) # y 노드의 대표 확인
        # 합치기 위해 각 대표를 비교해서 낮은값으로 대표로 선언
        if x_parent < y_parent: # x 대표가 낮아, y의 대표를 x로 선언
                graph[y_parent] = x_parent
        else:
                graph[x_parent] = y_parent # y 대표가 낮아, x의 대표를 y로 선언

의사코드

MakeSet(x):
    for each element x in X:
        parent[x] = x   // 각 원소의 부모를 자기 자신으로 설정
        rank[x] = 0     // 트리의 높이를 0으로 초기화

Find(x):
    if parent[x] ≠ x:
        parent[x] = Find(parent[x])  // 경로 압축: 부모를 대표로 업데이트
    return parent[x]

Union(x, y):
    rootX = Find(x)
    rootY = Find(y)

    if rootX == rootY:
        return  // 이미 같은 집합에 있음

    // Union by Rank: 높이가 낮은 트리를 높이가 높은 트리 아래에 붙인다.
    if rank[rootX] < rank[rootY]:
        parent[rootX] = rootY
    else if rank[rootX] > rank[rootY]:
        parent[rootY] = rootX
    else:
        parent[rootY] = rootX
        rank[rootX] += 1

프림(Prim) 알고리즘

한 정점에서 시작, 정점과 연결된 간선 중 가중치가 최소인 간선을 선택하면서 점진적으로 확장

예제

import heapq

def prim(graph, start=1):
    n = len(graph) - 1  # graph[0]은 사용하지 않는다고 가정
    visited = [False] * (n+1)
    min_heap = []
    # 시작 노드 방문 처리
    visited[start] = True
    for cost, neighbor in graph[start]:
        heapq.heappush(min_heap, (cost, start, neighbor))

    mst_cost = 0
    mst_edges = []

    while min_heap:
        cost, u, v = heapq.heappop(min_heap)
        if visited[v]:
            continue
        # 새로운 정점을 MST에 추가
        visited[v] = True
        mst_cost += cost
        mst_edges.append((u, v, cost))

        for next_cost, neighbor in graph[v]:
            if not visited[neighbor]:
                heapq.heappush(min_heap, (next_cost, v, neighbor))

    return mst_cost, mst_edges

# 사용 예제:
# 무방향 그래프를 인접 리스트로 표현
# graph[node] = [(비용, 인접노드), ...]
n = 6
graph = [[] for _ in range(n+1)]
edges = [
    (1, 2, 3),
    (1, 3, 2),
    (2, 3, 1),
    (2, 4, 4),
    (3, 4, 5),
    (3, 5, 6),
    (4, 5, 7),
    (4, 6, 8),
    (5, 6, 9)
]

# 양방향 그래프이므로 두 방향 모두 추가
for u, v, cost in edges:
    graph[u].append((cost, v))
    graph[v].append((cost, u))

mst_cost, mst_edges = prim(graph, start=1)
print("MST 총 비용:", mst_cost)
print("MST 간선들:")
for u, v, cost in mst_edges:
    print(u, "-", v, ":", cost)

흐름

1. 임의의 정점을 MST에 포함
2. 시작 정점과 연결된 모든 간선을 우선순위 큐에 넣음
3. 큐에서 최소 간선을 꺼내, 간선의 다른 쪽 정점이 아직 MST에 포함되지 않았따면 추가
4. 새 정점을 MST에 포함시킨 후, 정점과 연결된 간선들을 큐에 추가
5. 3~4를 반복