1922 네트워크 연결

Kruskal Algorithm

import sys


class Tree:
    def __init__(self, val, parent=None, height=1):
        self.val = val
        self.parent = parent
        self.height = height


def find(node):
    if node == node.parent:
        return node
    node.parent = find(node.parent)
    return node.parent


def union(a, b):
    pa = a.parent
    pb = b.parent

    if pa != pb:
        if pa.height < pb.height:
            pa.parent = pb
        else:
            pb.parent = pa

        if pa.height == pb.height:
            pa.height += 1


n = int(sys.stdin.readline().rstrip())
m = int(sys.stdin.readline().rstrip())
treedict = {}
graph = []

for _ in range(m):
    a, b, c = map(int, sys.stdin.readline().rstrip().split())

    if not(a in treedict.keys()):
        node = Tree(a)
        node.parent = node
        treedict[a] = node
    if not(b in treedict.keys()):
        node = Tree(b)
        node.parent = node
        treedict[b] = node

    graph.append((a, b, c))

graph.sort(key=lambda x: x[2])
dist = 0

for a, b, c in graph:
    pa = find(treedict[a])
    pb = find(treedict[b])

    if pa != pb:
        union(treedict[a], treedict[b])
        dist += c

sys.stdout.write(str(dist))

Prim Algorithm

import sys, heapq
from collections import defaultdict

n = int(sys.stdin.readline())
m = int(sys.stdin.readline())

graph = defaultdict(list)
visited = [0] * (n+1)

for _ in range(m):
    a, b, c = map(int, sys.stdin.readline().split())
    graph[a].append([c, a, b])
    graph[b].append([c, b, a])


def prim(graph, node):
    visited[node] = 1
    candidate = graph[node]
    heapq.heapify(candidate)
    total_weight = 0

    while candidate:
        w, u, v = heapq.heappop(candidate)

        if visited[v] == 0:
            visited[v] = 1
            total_weight += w

            for edge in graph[v]:
                if visited[edge[2]] == 0:
                    heapq.heappush(candidate, edge)

    return total_weight


sys.stdout.write(str(prim(graph, 1)))

1922 네트워크 연결