[BOJ] 10423 - 전기가 부족해 (Python)

 

 

아이디어

 

최소 스패닝 트리로 모든 도시에 전기를 공급할 수 있는 최소 비용을 구한다.

 

 

풀이 #1

 

프림 알고리즘으로 최소 스패닝 트리를 구현한다.

 

 

import sys
from heapq import heappush, heappop

input = sys.stdin.readline

INF = int(1e7)

def solution():
    N, M, K = map(int, input().split())
    plant = tuple(map(int, input().split()))
    graph = [dict() for _ in range(N+1)]
    for _ in range(M):
        u, v, w = map(int, input().split())
        graph[u][v] = graph[v][u] = min(graph[u].get(v, INF), w)
    MST = [0]*(N+1)
    D = [0]+[INF]*N
    heap = []
    for u in plant:
        MST[u] = 1
        D[u] = 0
        for v, w in graph[u].items():
            heappush(heap, (w, v))
    while heap:
        w, u = heappop(heap)
        if MST[u]:
            continue
        MST[u] = 1
        D[u] = w
        for v, _w in graph[u].items():
            if MST[v] == 0:
                heappush(heap, (_w, v))
    print(sum(D))

solution()

 

 

 

풀이 #2

 

크루스칼 알고리즘으로 최소 스패닝 트리를 구현한다.

 

import sys

input = sys.stdin.readline

INF = int(1e7)

def solution():

    def find(n):
        temp = n
        while n != group[n]:
            n = group[n]
        group[temp] = n
        return n

    def union(n1, n2):
        n1 = find(n1)
        n2 = find(n2)
        if n1 > n2:
            n1, n2 = n2, n1
        group[n2] = n1

    N, M, K = map(int, input().split())
    plant = tuple(map(int, input().split()))
    graph = [dict() for _ in range(N+1)]
    group = [i for i in range(N+1)]
    edges = []
    for _ in range(M):
        u, v, w = map(int, input().split())
        graph[u][v] = graph[v][u] = min(graph[u].get(v, INF), w)
        edges.append((w, u, v))
    edges.sort()
    for i in range(K):
        for j in range(i+1, K):
            union(plant[i], plant[j])
    res = cnt = 0
    for w, u, v in edges:
        if find(u) != find(v):
            union(u, v)
            res += w
            cnt += 1
            if cnt >= N:
                break
    print(res)

solution()