[백준] 물대기

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

1368번: 물대기

문제의 특징을 생각해봤을 때 논에 물을 끌어오려면 연결 그래프여야하고 각각의 연결 비용을 합쳤을 때 최소비용이 되어야하므로 자연스럽게 최소 신장 트리로 풀어야한다고 생각했습니다. 문제는 직접 논에 우물을 파는 것을 처리해야하는데 이 부분은 별도의 노드(코드에서는 other_node)를 만들어서 우물을 팔 때 발생하는 비용으로 각 노드와 연결시켰습니다. 이렇게 구축한 다음 프림 알고리즘을 돌리면 됩니다.

프림

import sys, heapq
from collections import defaultdict

input = sys.stdin.readline
INF = sys.maxsize
n = int(input())

wells = [-1] * (n + 1)


# 우물을 파는데 필요한 비용
for i in range(1, n + 1):
    wells[i] = int(input())

# 우물간의 인접 행렬
adj = [
    list(map(int, input().split()))
    for _ in range(n)
]

graph = defaultdict(list)


# 물을 길러올 별도의 노드
other_node = n + 1

def init():
    # 우물간의 경로를 설정
    for i in range(1, n + 1):
        for j in range(i + 1, n + 1):
            graph[i].append((j, adj[i - 1][j - 1]))
            graph[j].append((i, adj[i - 1][j - 1]))

    # 물을 길러오는 경우
    for i in range(1, n + 1):
        graph[other_node].append((i, wells[i]))
        graph[i].append((other_node, wells[i]))


def prim():
    # 1 ~ n + 1까지의 노드를 사용
    dist = [INF] * (n + 2)
    dist[other_node] = 0
    q = [(0, other_node)]
    visited = [False] * (n + 2)

    while q:
        d, v = heapq.heappop(q)

        if dist[v] < d:
            continue

        visited[v] = True

        for next_v, next_d in graph[v]:
            if not visited[next_v] and dist[next_v] > next_d:
                dist[next_v] = next_d
                heapq.heappush(q, (next_d, next_v))
    return sum(dist[1:])

init()
print(prim())

 

크루스칼

import sys, os

# sys.stdin = open(os.curdir + "\\input.txt", "r")

input = sys.stdin.readline

n = int(input())


# 우물을 파는데 드는 비용
well_price = [-1] * (n + 1)

edges = []

# 우물을 파야하는 노드
other_node = n + 1

# 우물에 해당하는 노드와 연결되어 있는 edge들을 넣음
for i in range(1, n + 1):
    price = int(input())
    edges.append((price, i, other_node))

# 인접 행렬 입력
adj = [[0] * (n + 1)] + [
    [0] + list(map(int, input().split()))
    for _ in range(n + 1)
]

# 인접 행렬의 edge들을 추가
for i in range(1, n + 1):
    for j in range(i + 1, n + 1):
        edges.append((adj[i][j], i, j))

edges.sort()
uf = [x for x in range(n + 2)]


def find(v):
    if uf[v] == v:
        return v

    uf[v] = find(uf[v])
    return uf[v]


answer = 0

for price, v, w in edges:
    # 부모 즉, 속해있는 집합을 찾는다.
    parent_w = find(w)
    parent_v = find(v)

    # 부모가 같은 경우는 이미 w, v가 MST에 속해있으므로 넘어간다.
    if parent_v == parent_w:
        continue

    answer += price

    if parent_v <= parent_w:
        uf[parent_w] = parent_v
    else:
        uf[parent_v] = parent_w

print(answer)