[백준] #1197 최소 스패닝 트리(python)

수영·2022년 8월 21일

MST Python3 백준

백준

목록 보기

42/117

📌문제

그래프가 주어졌을 때, 그 그래프의 최소 스패닝 트리를 구하는 프로그램을 작성하시오.

최소 스패닝 트리는, 주어진 그래프의 모든 정점들을 연결하는 부분 그래프 중에서 그 가중치의 합이 최소인 트리를 말한다.

입력

첫째 줄에 정점의 개수 V(1 ≤ V ≤ 10,000)와 간선의 개수 E(1 ≤ E ≤ 100,000)가 주어진다. 다음 E개의 줄에는 각 간선에 대한 정보를 나타내는 세 정수 A, B, C가 주어진다. 이는 A번 정점과 B번 정점이 가중치 C인 간선으로 연결되어 있다는 의미이다. C는 음수일 수도 있으며, 절댓값이 1,000,000을 넘지 않는다.

그래프의 정점은 1번부터 V번까지 번호가 매겨져 있고, 임의의 두 정점 사이에 경로가 있다. 최소 스패닝 트리의 가중치가 -2,147,483,648보다 크거나 같고, 2,147,483,647보다 작거나 같은 데이터만 입력으로 주어진다.

출력

첫째 줄에 최소 스패닝 트리의 가중치를 출력한다.

예제 입력

예제 출력

백준 1197번 문제

💡Idea

이 문제는 최소 스패닝 트리🎄를 만들기 위한 알고리즘을 그대로 구현하면 되는 문제입니다.

최소 스패닝 트리 알고리즘은 아래와 같이 두 가지가 있습니다.

크루스칼(Kruskal)
프림(Prim)

두 알고리즘에 대해서 더 알고 싶다면 여기를 클릭해주세요!

💻코드1

⏰ 시간 : 388 ms / 메모리 : 51616 KB

import sys
input = sys.stdin.readline

V, E = map(int, input().split())
tree = []
for _ in range(E):
    A, B, C = map(int, input().split()) # A에서 B로 가는 간선의 weight C
    tree.append([A, B, C])

def find(mst, v): # v가 속한 집합의 대표를 찾는 함수
    if mst[v] != v: 
        mst[v] = find(mst, mst[v]) # 집합의 대표를 찾을 때까지 재귀
    return mst[v]

def union(mst, u, v): # u와 v가 속한 집합 두 개를 union
    up, vp = find(mst, u), find(mst, v) 
    if up < vp : # 집합의 대표가 더 작은 값이 합친 집합의 대표
        mst[vp] = up
    else:
        mst[up] = vp

def kruskal(tree):
    total = 0
    mst = [x for x in range(V+1)] # 처음에는 각 정점이 집합의 대표
    tree.sort(key=lambda x : x[2]) # weight를 기준으로 정렬

    for i in range(E): # weight가 작은 간선부터 확인
        a, b, c = tree[i]
        if find(mst, a) != find(mst, b): # 두 정점이 서로 다른 집합에 속한 경우
            union(mst, a, b) # a와 b의 집합을 union
            total += c # a와 b를 연결하는 간선의 weight를 MST 가중치에 추가
    return total

print(kruskal(tree))

📝코드 설명1

크루스칼 알고리즘을 구현하여 문제를 해결한 코드입니다.

[정점1, 정점2, 간선의 weight]으로 각 간선들을 리스트 형태로 저장하여 사용하였습니다.

크루스칼은 그리디 기법을 활용하여 가장 비용이 적은 간선부터 차례로 선택하기 때문에, 먼저 간선의 비용을 기준으로 간선들을 정렬해줍니다.

그리고 비용이 적은 간선부터 하나씩 확인하면서 해당 간선의 두 정점이 이미 서로 같은 집합에 있어 사이클을 형성하지 않는가를 확인하고, 만약 그렇지 않다면 해당 간선을 MST로 선택하여 두 정점의 집합을 합쳐줍니다. (이 때, 총 비용을 계산하기 위하여 간선의 비용을 total에 더해줍니다.)

정점의 사이클 형성 여부 확인과, 두 정점을 같은 집합으로 합치는 연산을 위하여 유니온 파인드(Union-Find, Disjoint Set)을 사용하였습니다.

💻코드2

⏰ 시간 : 652 ms / 메모리 : 68936 KB

import heapq
import sys
input = sys.stdin.readline

V, E = map(int, input().split())
tree = {x: [] for x in range(V+1)}
for _ in range(E):
    A, B, C = map(int, input().split()) # A에서 B로 가는 간선의 weight C
    tree[A].append([C, B])
    tree[B].append([C, A])

def prim(tree, start):
    total = 0
    mst = set([start]) # 현재 MST에 포함된 노드 집합
    candidate = tree[start] # 시작 정점과 연결된 [노드, 간선 weight] 리스트
    heapq.heapify(candidate) # 우선순위 큐로 만들어 줌
    while candidate :
        c, b = heapq.heappop(candidate) # 가장 낮은 비용의 간선 pop
        if b not in mst: # 아직 MST에 포함되지 않은 정점의 경우
            mst.add(b) # MST에 정점 추가
            total += c # 총 비용에 weight 추가
            for w, n in tree[b]: # 정점의 간선 리스트 확인
                if n not in mst: 
                	# 아직 MST에 포함되지 않은 정점의 경우 우선순위 큐 추가
                    heapq.heappush(candidate, [w, n])
    return total

print(prim(tree, 1))