[Python/파이썬] 백준 10423번 전기가 부족해

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

코드

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

n, m, k = map(int, input().split())
power_stations = set(map(int, input().split()))
parent = [i for i in range(n+1)]


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


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

    if a == b:
        return False

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


cables = []
for _ in range(m):
    u, v, w = map(int, input().split())
    heappush(cables, (w, u, v))

answer = 0
while cables:
    cost, u, v = heappop(cables)
    if find_parent(u) in power_stations and find_parent(v) in power_stations:
        continue

    if union(u, v):
        answer += cost

    is_all_connected = True
    for i in range(1, n+1):
        if find_parent(i) not in power_stations:
            is_all_connected = False
            break
    if is_all_connected:
        break

print(answer)

설명

크루스칼(Kruskal) 알고리즘을 사용하여 문제를 풀었다. 

크루스칼은 가장 가중치가 작은 간선을 하나씩 선택하며 MST를 만드는 알고리즘이다. 이 알고리즘을 수행할 때는 2개의 함수가 필요하다.

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


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

    if a == b:
        return False

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

- `find_parent(x)`: 인자로 받은 x가 속한 그래프의 루트 노드를 찾는 함수.

- `union(a, b)`: 인자로 받은 a, b가 만약 사이클을 생성하지 않는다면 두 정점을 같은 그래프로 합치는 함수. 여기서 조금 다른 점은 합칠 때 발전소가 있는 정점이 루트가 될 수 있도록 한다는 점이다. 각 정점이 속한 그래프에 이미 발전소가 속해 있다면 이 문제는 하나의 MST를 만드는 것이 아니라 발전소가 있는 정점이 하나만 포함되는 여러 개의 MST를 만드는 문제이기 때문에 이렇게 구현했다.

 

이렇게 필요한 함수들을 구현해둔 뒤 다음 과정을 진행한다.

cables = []
for _ in range(m):
    u, v, w = map(int, input().split())
    heappush(cables, (w, u, v))

입력으로 받은 케이블은 최소힙에 넣어서 가중치가 작은 간선부터 union할 수 있도록 한다.

answer = 0
while cables:
    cost, u, v = heappop(cables)
    if find_parent(u) in power_stations and find_parent(v) in power_stations:
        continue

    if union(u, v):
        answer += cost

    is_all_connected = True
    for i in range(1, n+1):
        if find_parent(i) not in power_stations:
            is_all_connected = False
            break
    if is_all_connected:
        break

print(answer)

1. 최소힙 cables에서 `heappop`으로 가장 가중치가 작은 간선을 뽑는다.

2. 뽑은 간선의 양 끝 정점이 발전소와 연결된 정점인지 확인한다. 둘 다 이미 발전소와 연결되어 있다면 이후 과정은 진행할 필요 없다.

3. 두 정점을 `union`한다. 이 때 사이클을 만들지 않는 경우에만 이 간선을 사용하게 되므로 `union()`이 `True`를 반환하는 경우에만 `answer`에 가중치를 더해준다.

4. 모든 정점이 발전소와 연결되었는지 확인한다. `parent[1:]`의 모든 원소의 루트가 `power_stations`에 속하는 값인지 확인하면 된다.