Coding 공부

import sys
import heapq

input = sys.stdin.readline

N,M,K = map(int,input().split())

graph = [{} for i in range(N+1)]
plant = list(map(int,input().split()))
INF = float('inf')
for _ in range(M):
    u,v,w = map(int,input().split())
    graph[u][v] = min(graph[u].get(v,float('inf')),w)
    graph[v][u] = min(graph[v].get(u,float('inf')),w)



MST_DISTANCE = [INF for i in range(N+1)]
visited = [True]*(N+1)

result = 0
node_list = []
for start in plant:
    heapq.heappush(node_list,(0,start))
    MST_DISTANCE[start] = 0

while node_list:
    dis,node = heapq.heappop(node_list)
    if not visited[node]:
        continue
    result += dis
    visited[node] = False
    for next_node in graph[node]:
        if MST_DISTANCE[next_node] >graph[node][next_node]:
            MST_DISTANCE[next_node] = graph[node][next_node]
            heapq.heappush(node_list,(MST_DISTANCE[next_node],next_node))


print(result)

서로 다른 발전소끼리 연결이 되면 안되므로, 프림알고리즘을 하는데, 발전소의 위치를 전부 0으로 초기화 해주고,

전부 heapq에 넣어주고, Prim 알고리즘을 돌려주면 된다.

import sys
input = sys.stdin.readline

def union(a,b):
    A = find_parent(a)
    B = find_parent(b)
    if A != B:
        if rank[A] < rank[B]:
            A,B = B,A
        make_set[B] = A
        if rank[A] == rank[B]:
            rank[A] += 1

def find_parent(ind):
    if make_set[ind] == ind:
        return ind
    else:
        make_set[ind] = find_parent(make_set[ind])
        return make_set[ind]


N,M,K = map(int,input().split())


plant = list(map(int,input().split()))
make_set = [i for i in range(N+1)]
rank = [1 for _ in range(N+1)]
edges = []
for _ in range(M):
    u,v,w = map(int,input().split())
    edges.append((u,v,w))

edges.sort(key=lambda x : -x[2])


for k in range(1,K):
    union(plant[k],plant[k-1])

cnt = 1
result = 0
while cnt <(N-K+1):
    x,y,weight = edges.pop()
    if find_parent(x) != find_parent(y):
        union(x,y)
        result += weight
        cnt += 1

print(result)

    두번째 풀이는 크루스칼을 이용해서 풀었다.

발전소가 서로 연결이 되면 안되므로, 처음부터 모든 발전소를 하나의 union으로 merge를 해준다.

그리고 난뒤에 크루스칼 알고리즘을 해주면 된다.

그리고 전체 간선 수는 전체 노드의 수 - 1 이지만, 여기서는 (N-1)-(K-1)의 수 만큼만 해주면 된다.

이 문제를 처음에 어렵게 생각해서 각 플랜트만의 프림알고리즘을 돌려서, 최저값을 갱신하는 식으로 했다.

하지만, 같이 연결이 안되기만 하면 되니깐 프림 알고리즘을 하면서, 처음 스타트 위치를 전부 넣어주기만 하면 됬다.

이 문제에 더 어울리는 알고리즘은 프림보다, 크루스칼 알고리즘인 것 같다.