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

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

 

프림 알고리즘으로 해결하기 적절한 문제이다.

# 백준 10423 전기가 부족해
# 문제 요약:
#   - N개의 도시, M개의 도로, K개의 발전소가 있음
#   - 발전소가 있는 도시부터 전기를 연결하면서 모든 도시에 전기 공급
#   - 최소 비용으로 모든 도시에 전기 연결 (최소 신장 트리 문제, 다중 시작점)

import sys
import heapq
input = sys.stdin.readline

# 입력
N, M, K = map(int, input().split())
start_citys = list(map(int, input().split()))  # 발전소가 있는 도시 번호

ans = 0             # 최종 최소 비용
cnt = 0             # MST에 포함된 도시 수
visited = [False for _ in range(N+1)]  # 도시 방문 여부
conn = [[] for _ in range(N+1)]        # 도시 연결 정보 adjacency list
curr_costs = []                         # 현재 선택 가능한 간선 최소 힙

# 도로 정보 입력
for _ in range(M):
    A, B, cost = map(int, input().split())
    conn[A].append((B, cost))
    conn[B].append((A, cost))

# 발전소 도시들을 MST 초기 노드로 설정
while start_citys:
    start_city = start_citys.pop()   # 발전소 도시
    visited[start_city] = True       # MST에 포함
    # 발전소 도시와 연결된 간선을 힙에 추가
    for conn_city, cost in conn[start_city]:
        heapq.heappush(curr_costs, [cost, conn_city])
    cnt += 1  # MST 포함 도시 수 증가

# Prim 알고리즘 수행
while curr_costs:
    cost, next_city = heapq.heappop(curr_costs)  # 가장 작은 비용 간선 선택
    if not visited[next_city]:
        ans += cost                  # MST 비용 누적
        visited[next_city] = True    # MST에 도시 포함

        # 새로 포함된 도시와 연결된 간선 힙에 추가
        for next_next_city, next_cost in conn[next_city]:
            heapq.heappush(curr_costs, [next_cost, next_next_city])
        cnt += 1

    if cnt == N:   # 모든 도시가 MST에 포함되면 종료
        break

print(ans)  # 최소 비용 출력

프림알고리즘으로 해결하였다.

heapq와 visited를 이용하여 현 시점에서 가장 cost가 가장 낮고, 사이클이 생기지 않도록 선을 연결한다.

 

GPT가 제시한 크루스칼 알고리즘으로 해결한 풀이도 첨부한다.

 

# 백준 10423 전기가 부족해 - Kruskal 알고리즘 버전
import sys
input = sys.stdin.readline

# Union-Find 초기화
def find(x):
    if parent[x] != x:
        parent[x] = find(parent[x])
    return parent[x]

def union(a, b):
    a = find(a)
    b = find(b)
    if a != b:
        parent[a] = b
        return True
    return False

# 입력
N, M, K = map(int, input().split())
start_citys = list(map(int, input().split()))  # 발전소 도시 번호

parent = [i for i in range(N+1)]

# 발전소 도시를 하나의 가상 집합으로 초기화
# 즉, 모든 발전소는 이미 연결된 상태라고 가정
for i in start_citys:
    parent[i] = start_citys[0]

# 모든 간선 입력
edges = []
for _ in range(M):
    A, B, cost = map(int, input().split())
    edges.append((cost, A, B))

# 비용 기준 오름차순 정렬
edges.sort()

ans = 0
for cost, A, B in edges:
    # A, B를 union할 수 있으면 MST에 추가
    if union(A, B):
        ans += cost

print(ans)

크루스칼 알고리즘에서는 발전소들을 가상으로 연결하고 시작하였다.

 

 

MST가 조금씩 익숙해진다.

크루스칼 - UNION-FIND

프림 - HEAPQ, visited