크루스칼(Kruskal) 알고리즘 구현 (python)

단간단간·2024년 4월 24일

크루스칼

알고리즘 유형

목록 보기

2/3

최소 신장 트리 (Minimum Spanning Tree, MST)

주어진 가중치 그래프의 모든 정점을 최소한의 비용으로 연결하는 부분 그래프를 의미한다.
주로 크루스칼(Kruskal) 알고리즘 또는 프림(Prim) 알고리즘을 사용해서 해결할 수 있다.

크루스칼 알고리즘 (`Kruskal`)

가장 가벼운 가중치를 가진 간선부터 선택하면서 MST를 구성한다.
최종적으로 선택될 간선의 총 개수는 (총 노드개수 - 1) 이다.
간선을 선택할 때는, 사이클이 형성되지 않도록 주의한다. 사이클 형성 여부를 확인하기 위해 Union-Find 자료구조를 사용한다.

python

# 특정 원소가 속한 집합 찾기
def find_parent(parent: list, x: int):
    # 루트 노드 찾을 때까지 재귀 호출
    if parent[x] != x:
        parent[x] = find_parent(parent, parent[x])

    return parent[x]


# 두 원소가 속한 집합 합치기
def union_parent(parent, a, b):
    a_parent = find_parent(parent, a)
    b_parent = find_parent(parent, b)

    if a_parent < b_parent:
        parent[b_parent] = a_parent
    else:
        parent[a_parent] = b_parent

    return parent


# v: 노드 개수, edges: 간선에 대한 정보(비용, 연결된노드1, 연결된노드2)
def solution(v: int, edges: list):
    # 부모 테이블 자기자신으로 초기화 (사이클 여부 확인 목적)
    parent = [i for i in range(v + 1)]

    # 간선을 비용순으로 정렬
    edges.sort()

    result = 0
    count = 0
    for cost, a, b in edges:
        # 사이클이 발생하지 않는 경우에만 집합에 포함
        if find_parent(parent, a) != find_parent(parent, b):
            parent = union_parent(parent, a, b)
            result += cost
            count += 1

        # 간선 숫자 = 노드 개수 - 1
        if count == v - 1:
            break

    return result


if __name__ == "__main__":
    result = solution(
        v=5,
        edges=[(1, 1, 2), (2, 3, 4), (3, 2, 4), (4, 1, 3), (5, 3, 5), (2, 2, 3)]
    )

    print(result)