[Python] Kruskal Algorithm (크루스칼 알고리즘)

크루스칼 알고리즘의 이론은 https://brownbears.tistory.com/396에 설명이 되어 있습니다. 아래는 파이썬으로 기본적인 크루스칼 알고리즘을 어떻게 구현했는지에 대한 내용입니다.

코드에 앞서 크루스칼 알고리즘의 중요한 점은 아래와 같습니다.

1. edge가 가장 작은 순으로 선택해 나감
2. edge 연결 후, cycle이 생기면 안됨

즉, 가장 작은 edge를 선택할 때마다, 그동안 연결한 vertices끼리 cycle이 생기는지 확인을 해야 합니다. 이러한 cycle이 생겼는지 확인하는 것은 union-find (disjoint-set) 알고리즘을 사용합니다. union-find (disjoint-set) 알고리즘의 설명은 https://brownbears.tistory.com/460에 정리되어 있습니다.

아래는 kruskal 알고리즘 코드의 예시입니다.

class DisjointSet: def __init__(self, n): self.parent = {} self.rank = {} for i in range(n): self.parent[i] = i self.rank[i] = 0 def find(self, v): if self.parent[v] != v: self.parent[v] = self.find(self.parent[v]) return self.parent[v] def union(self, root1, root2): # 연결 정보를 갱신할 때는 작은 값을 기준으로 갱신 if self.rank[root1] > self.rank[root2]: self.parent[root2] = root1 else: self.parent[root1] = root2 # self.rank[root2]가 self.rank[root1]보다 크다고 봤으므로 두 값이 같은 경우, self.rank[root2]의 값을 1 올려준다. if self.rank[root1] == self.rank[root2]: self.rank[root2] += 1 def kruskal(n, info): minimum_weight = 0 disjointset = DisjointSet(n) result = [] for data in sorted(info, key=lambda cost: cost[2]): v, u, weight = data root1 = disjointset.find(v) root2 = disjointset.find(u) # root1과 root2가 같은데 연결을 하면 사이클이 생김 if root1 != root2: disjointset.union(root1, root2) result.append(data) minimum_weight += data[2] return result, minimum_weight kruskal(6, [[0, 1, 5], [0, 3, 2], [0, 4, 3], [1, 4, 1], [3, 4, 10], [1, 2, 2], [2, 5, 3], [4, 5, 4]]) # 결과 [[1, 4, 1], [0, 3, 2], [1, 2, 2], [0, 4, 3], [2, 5, 3]], 11

위 코드 kruskal() 함수에서 첫 번째 인자은 n은 vertices의 총 개수이고 info는 각 vertex 간 연결 정보와 link의 weight를 나타냅니다. 즉 예시에서 6은 vertex가 6개 있다는 의미이며 두 번째 인자에서 [0, 1, 5] 은 vertex 0과 1이 연결되어 있고 이 link의 weight은 5라는 의미입니다.

이러한 예시를 진행하면 최소 weight는 11이 나오게 됩니다.