[Python/파이썬] 백준 16202번 MST 게임

https://www.acmicpc.net/problem/16202

 

코드

from heapq import heappush
import sys
input = sys.stdin.readline

n, m, k = map(int, input().split())
graph = []
for i in range(m):
    x, y = map(int, input().split())
    graph.append((i+1, x, y))


def find_parent(x, parent):
    if x != parent[x]:
        parent[x] = find_parent(parent[x], parent)
    return parent[x]


def union(a, b, parent):
    a = find_parent(a, parent)
    b = find_parent(b, parent)

    if a == b:
        return False

    if a < b:
        parent[b] = a
    else:
        parent[a] = b
    return True


possible = [True]*(m+1)
answer = [0]*k
for i in range(k):
    parent = [j for j in range(n+1)]
    # Kruskal
    edges = []
    total_score = 0
    for cost, x, y in graph:
        if possible[cost] and find_parent(x, parent) != find_parent(y, parent):
            union(x, y, parent)
            heappush(edges, (cost, x, y))
            total_score += cost
            if len(edges) == n-1:
                break
    possible[edges[0][0]] = False
    if len(edges) == n-1:
        answer[i] = total_score
    else:
        break
print(*answer)

 

일단 문제가 최소 스패닝 트리(MST)를 만들어야 하니 선택한 간선이 사이클을 이루는지 판별할 find_parent 함수와 트리에 간선을 추가할 union 함수를 구현한다.

 

그리고 크루스칼 알고리즘을 이용해서 사용 가능한 간선들로 MST를 만드는 과정을 k번 반복한다.

1. 원래는 간선들 중 가장 작은 가중치를 갖는 간선을 하나씩 선택해야 하는데 애초에 입력받을 때 입력 순서가 가중치가 됐기 때문에 graph에 저장된 그대로 사용해도 된다. 다만 k번 반복하며 MST에 포함된 간선 중 가장 작은 간선을 하나씩 빼기 때문에 이를 체크하기 위해&선택된 간선의 개수를 체크하기 위해 edges라는 최소힙에 선택한 간선들을 넣었다.
2. MST를 만드는 과정은 간선을 n-1개 고르면 끝나고 가장 작은 간선은 이제 사용할 수 없게 되므로 각 간선의 사용 여부를 표시하는 리스트인 possible에 edges의 첫번째 원소로 있는 간선을 False로 바꿔준다.
3. 간선을 n-1개 고른다면 answer[i]=total_score를 해주면 되지만, 고르지 못한다면 더이상 MST를 만들 수 없다는 뜻이므로 반복문을 탈출한다.