"""Prim's Algorithm.⏎
⏎
    Determines the minimum spanning tree(MST) of a graph using the Prim's Algorithm.⏎
⏎
    Details: https://en.wikipedia.org/wiki/Prim%27s_algorithm⏎
"""⏎
⏎
import heapq as hq⏎
import math⏎
from typing import Iterator⏎
⏎
⏎
class Vertex:⏎
    """Class Vertex."""⏎
⏎
    def __init__(self, id):⏎
        """⏎
        Arguments:⏎
            id - input an id to identify the vertex⏎
        Attributes:⏎
            neighbors - a list of the vertices it is linked to⏎
            edges     - a dict to store the edges's weight⏎
        """⏎
        self.id = str(id)⏎
        self.key = None⏎
        self.pi = None⏎
        self.neighbors = []⏎
        self.edges = {}  # {vertex:distance}⏎
⏎
    def __lt__(self, other):⏎
        """Comparison rule to < operator."""⏎
        return self.key < other.key⏎
⏎
    def __repr__(self):⏎
        """Return the vertex id."""⏎
        return self.id⏎
⏎
    def add_neighbor(self, vertex):⏎
        """Add a pointer to a vertex at neighbor's list."""⏎
        self.neighbors.append(vertex)⏎
⏎
    def add_edge(self, vertex, weight):⏎
        """Destination vertex and weight."""⏎
        self.edges[vertex.id] = weight⏎
⏎
⏎
def connect(graph, a, b, edge):⏎
    # add the neighbors:⏎
    graph[a - 1].add_neighbor(graph[b - 1])⏎
    graph[b - 1].add_neighbor(graph[a - 1])⏎
    # add the edges:⏎
    graph[a - 1].add_edge(graph[b - 1], edge)⏎
    graph[b - 1].add_edge(graph[a - 1], edge)⏎
⏎
⏎
def prim(graph: list, root: Vertex) -> list:⏎
    """Prim's Algorithm.⏎
⏎
        Runtime:⏎
            O(mn) with `m` edges and `n` vertices⏎
⏎
        Return:⏎
            List with the edges of a Minimum Spanning Tree⏎
⏎
        Usage:⏎
            prim(graph, graph[0])⏎
    """⏎
    a = []⏎
    for u in graph:⏎
        u.key = math.inf⏎
        u.pi = None⏎
    root.key = 0⏎
    q = graph[:]⏎
    while q:⏎
        u = min(q)⏎
        q.remove(u)⏎
        for v in u.neighbors:⏎
            if (v in q) and (u.edges[v.id] < v.key):⏎
                v.pi = u⏎
                v.key = u.edges[v.id]⏎
    for i in range(1, len(graph)):⏎
        a.append((int(graph[i].id) + 1, int(graph[i].pi.id) + 1))⏎
    return a⏎
⏎
⏎
def prim_heap(graph: list, root: Vertex) -> Iterator[tuple]:⏎
    """Prim's Algorithm with min heap.⏎
⏎
        Runtime:⏎
            O((m + n)log n) with `m` edges and `n` vertices⏎
⏎
        Yield:⏎
            Edges of a Minimum Spanning Tree⏎
⏎
        Usage:⏎
            prim(graph, graph[0])⏎
    """⏎
    for u in graph:⏎
        u.key = math.inf⏎
        u.pi = None⏎
    root.key = 0⏎
⏎
    h = [v for v in graph]⏎
    hq.heapify(h)⏎
⏎
    while h:⏎
        u = hq.heappop(h)⏎
        for v in u.neighbors:⏎
            if (v in h) and (u.edges[v.id] < v.key):⏎
                v.pi = u⏎
                v.key = u.edges[v.id]⏎
                hq.heapify(h)⏎
⏎
    for i in range(1, len(graph)):⏎
        yield (int(graph[i].id) + 1, int(graph[i].pi.id) + 1)⏎
⏎
⏎
def test_vector() -> None:⏎
    """⏎
    # Creates a list to store x vertices.⏎
    >>> x = 5⏎
    >>> G = [Vertex(n) for n in range(x)]⏎
⏎
    >>> connect(G, 1, 2, 15)⏎
    >>> connect(G, 1, 3, 12)⏎
    >>> connect(G, 2, 4, 13)⏎
    >>> connect(G, 2, 5, 5)⏎
    >>> connect(G, 3, 2, 6)⏎
    >>> connect(G, 3, 4, 6)⏎
    >>> connect(G, 0, 0, 0)  # Generate the minimum spanning tree:⏎
    >>> G_heap = G[:]⏎
    >>> MST = prim(G, G[0])⏎
    >>> MST_heap = prim_heap(G, G[0])⏎
    >>> for i in MST:⏎
    ...     print(i)⏎
    (2, 3)⏎
    (3, 1)⏎
    (4, 3)⏎
    (5, 2)⏎
    >>> for i in MST_heap:⏎
    ...     print(i)⏎
    (2, 3)⏎
    (3, 1)⏎
    (4, 3)⏎
    (5, 2)⏎
    """⏎
⏎
⏎
if __name__ == "__main__":⏎
    import doctest⏎
⏎
    doctest.testmod()⏎