import math
from min_heap import Element, MinHeap

class Edge:
    def __init__(self, u, v, w):
        self.u = u
        self.v = v
        self.w = w

class Graph:
    def __init__(self, vertex_num):
        self.adj_list = [[] for _ in range(vertex_num)]
        self.edge_list = []

        self.vertex_num = vertex_num

    def add_edge(self, u, v, w):

➊

       self.adj_list[u].append((v, w))
       self.adj_list[v].append((u, w))

        self.edge_list.append(Edge(u, v, w))

    def MST_prim(self):
        mst = Graph(self.vertex_num) ➋

        w_list = [math.inf for _ in range(self.vertex_num)]  ➌

        TV = set() ➍
        h = MinHeap() ➎

        for i in range(1, self.vertex_num):
            h.push(Element(i, math.inf, None)) ➏

        w_list[0] = 0 ➐
        h.push(Element(0, 0, None))

        while not h.is_empty():
            elem_v = h.pop() ➑
            v = elem_v.v
            w = elem_v.w
            _from = elem_v._from

            TV.add(v) ➒

➓

           if _from != None:
               mst.add_edge(v, _from, w)

            adj_v = self.adj_list[v] 11
            for u, w_u_v in adj_v: 12
                if u not in TV and w_u_v < w_list[u]:
                    w_list[u] = w_u_v
                    h.decrease_weight(Element(u, w_u_v, v))
        return mst

➊ (정점, 에지의 가중치)를 인접 리스트에 추가

➋ 최종적으로 만들 MST

➌ w_list: 각 정점의 w 값을 담아 두는 배열

➍ TV = {}: MST 정점의 집합, 시작 노드부터 하나씩 채워 나갑니다.

➎ min heap에 w와 from을 가진 정점을 담아 둡니다. heap 초기화: w -> inf, from -> None

➏ Element(v, w, _from)

➐ 시작 노드인 0은 w -> 0, from -> None

➑ 가중치가 가장 작은 에지 (from, v): w 정보를 가진 정점 Element v

➒ TV에 정점을 추가

➓ TE에 에지 추가

11 TV에 정점이 추가되면 인접 정점 중 트리 밖에 있는 정점에 대해 업데이트 시도

adj[v] = 정점 v에 인접한 정점 집합

12 w_u_v: weight(u, v)