[Python] 백준 16398 - 행성 연결 문제 풀이

Overview

BOJ 16398번 행성 연결 Python 문제 풀이
분류: Minimum Spanning Tree (최소 스패닝 트리), Kruskal Algorithm (크루스칼 알고리즘)

---

문제 페이지

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

---

풀이 코드

from sys import stdin, setrecursionlimit
from typing import List
setrecursionlimit(10 ** 9)


parents = []


def find(x: int) -> int:
    if parents[x] == x:
        return x
    parents[x] = find(parents[x])
    return parents[x]


def union(x: int, y: int) -> None:
    x, y = find(x), find(y)
    if x < y:
        parents[y] = x
    else:
        parents[x] = y


def kruskal(n: int, planets: List[List[int]]) -> int:
    edges = []

    for i in range(n):
        for j in range(i + 1, n):
            edges.append((planets[i][j], i, j))

    edges.sort()
    cost = 0
    for edge, a, b in edges:
        if find(a) == find(b):
            continue
        union(a, b)
        cost += edge

    return cost


def main():
    def input():
        return stdin.readline().rstrip()
    global parents

    n = int(input())
    planets = [list(map(int, input().split())) for _ in range(n)]
    parents = list(i for i in range(n))
    print(kruskal(n, planets))


if __name__ == "__main__":
    main()

크루스칼 알고리즘을 이용하여, 최소 신장 트리를 만들어 답을 구한다.