백준 1833 - 고속철도 설계하기 (Python)

문제 정보

• https://www.acmicpc.net/problem/1833
• MST

---

문제 정리

N개의 도시가 있고, 각 도시 사이에는 고속철도가 이미 깔려있는 곳이 있고 깔리지 않은 곳이 있다. 여기에 임의의 도시에서 다른 임의의 도시로 이동할 수 있게 고속철도를 추가로 설치하려고 한다. 고속철도 설치에 필요한 최소 비용을 구하는 문제이다.

---

풀이

"임의의 도시에서 다른 임의의 도시로 이동" 이 말은 트리를 구하라는 말이고. 간선을 적절하게 선택하여 최소 비용을 구하는 것이 목적이므로 해당 그래프에서 MST를 구하면 된다.
일단, 이미 깔려있는 고속철도는 지출된 비용이므로 해당 간선들의 비용은 총 비용에 먼저 더해준다. 그 뒤, 이 간선들의 비용을 0으로 바꾸어 준다. 그 뒤, 전체 간선에 대해 MST를 구하고, MST를 구성하는 간선의 비용을 총 비용에 더해주면 문제에서 구하는 값을 구할 수 있다. 추가로 설치해야 하는 간선도 구해야 하긴 하는데, 이건 그냥 MST를 구성하면서 비용이 0이 아닌 간선들을 리스트에 저장해뒀다가 출력하면 된다.

그냥 저냥 무난한 MST 문제인듯?

---

코드

# 백준 1833

import io

input = io.BufferedReader(io.FileIO(0), 1<<18).readline


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


def union(parent, a, b):
    a = find(parent, a)
    b = find(parent, b)
    parent[max(a, b)] = min(a, b)


def Kruskal(V, edge, prevCost):
    total = 0
    edgeMST = []
    parent = [i for i in range(V+1)]
    edge.sort(key= lambda x: x[2])

    for cur in edge:
        startV, endV, cost = cur
        if find(parent, startV) != find(parent, endV):
            union(parent, startV, endV)
            if cost != 0:
                total += cost
                edgeMST.append((startV, endV))

    return [[total+prevCost, len(edgeMST)]], edgeMST


def main():
    N = int(input())
    edge = []
    prevCost = 0
    for s in range(1, N+1):
        A = list(map(int, input().split()))
        for i in range(s, N):
            if A[i] > 0:
                edge.append((s, i+1, A[i]))
            elif A[i] < 0:
                edge.append((s, i+1, 0))
                prevCost -= A[i]
    for i in Kruskal(N, edge, prevCost):
        for j in i:
            print(*j)


main()