prim_algorithm.py

from typing import List, Tuple, Set
from abc import ABC, abstractmethod


class Vertex:
    def __init__(self, key: str, data: str = None) -> None:
        self.key = key
        self.data = data
        self.col = -1

    def __eq__(self, other) -> bool:
        return self.key == other.key

    def __hash__(self) -> int:
        return hash(self.key)

    def __repr__(self) -> str:
        return f"{self.key}"


class Graph(ABC):
    @abstractmethod
    def is_empty(self) -> bool:
        pass

    @abstractmethod
    def size(self) -> int:
        pass

    @abstractmethod
    def insert_vertex(self, vertex: Vertex) -> None:
        pass

    @abstractmethod
    def insert_edge(self, vertex1: Vertex, vertex2: Vertex, edge: int = 1) -> None:
        pass

    @abstractmethod
    def delete_vertex(self, vertex: Vertex) -> None:
        pass

    @abstractmethod
    def delete_edge(self, vertex1: Vertex, vertex2: Vertex) -> None:
        pass

    @abstractmethod
    def edges(self) -> List[Tuple[Vertex, Vertex]]:
        pass

    @abstractmethod
    def color(self, vertex: Vertex, t: str) -> None:
        pass

    @abstractmethod
    def prim(self, v: Vertex):
        pass


class GraphList(Graph):
    def __init__(self) -> None:
        self.adj_list = {}
        self.list_vertex = []
        self.map_vertex = {}
        self.in_tree = set()
        self.distance = {}
        self.parent = {}

    def is_empty(self) -> bool:
        return not self.list_vertex

    def size(self) -> int:
        if self.is_empty():
            return 0
        return len(self.list_vertex)

    def insert_vertex(self, vertex: Vertex) -> None:
        if vertex not in self.adj_list:
            self.adj_list[vertex] = {}
            self.list_vertex.append(vertex)
            self.map_vertex[vertex] = self.size() - 1

    def insert_edge(self, vertex1: Vertex, vertex2: Vertex, edge: int = 1) -> None:
        if vertex1 not in self.adj_list:
            self.insert_vertex(vertex1)
        if vertex2 not in self.adj_list:
            self.insert_vertex(vertex2)

        self.adj_list[vertex1][vertex2] = edge
        self.adj_list[vertex2][vertex1] = edge

    def delete_vertex(self, vertex: Vertex) -> None:
        if vertex in self.adj_list:
            ind = self.map_vertex[vertex]
            del self.adj_list[vertex]
            del self.map_vertex[vertex]
            del self.list_vertex[ind]
            for v, i in self.map_vertex.items():
                if i > ind:
                    self.map_vertex[v] -= 1
            for current_vertex, neighs in self.adj_list.items():
                if vertex in neighs:
                    del self.adj_list[current_vertex][vertex]

    def delete_edge(self, vertex1: Vertex, vertex2: Vertex) -> None:
        if vertex1 in self.adj_list and vertex2 in self.adj_list:
            if vertex1 in self.adj_list[vertex2] and vertex2 in self.adj_list[vertex1]:
                del self.adj_list[vertex1][vertex2]
                del self.adj_list[vertex2][vertex1]

    def edges(self) -> List[Tuple[Vertex, Vertex]]:
        if not self.is_empty():
            l = []
            for current_vertex, neighs in self.adj_list.items():
                for neighbour_vertex in neighs:
                    tup = (current_vertex.key, neighbour_vertex.key)
                    l.append(tup)
            return l

    def prim(self, v: Vertex):
        mst = GraphList()
        length = 0
        for ver in self.list_vertex:
            self.distance[ver] = float('inf')
            self.parent[ver] = -1
        while v not in self.in_tree:
            self.in_tree.add(v)
            for neigh in self.adj_list[v]:
                if neigh not in self.in_tree and \
                        self.adj_list[v][neigh] < self.distance[neigh]:
                    self.distance[neigh] = self.adj_list[v][neigh]
                    self.parent[neigh] = v
            min_dist = float('inf')
            min_ver = v
            for to_ver, dist in self.distance.items():
                if to_ver not in self.in_tree and dist < min_dist:
                    min_ver = to_ver
                    min_dist = dist
            mst.insert_edge(self.parent[min_ver], min_ver, self.distance[min_ver])
            v = min_ver
        return mst, length

    def __check_color_availability(self, v: Vertex, visited: Set[Vertex], cols: List[int]) -> int:
        cols_temp = set()
        for vertex in self.list_vertex:
            if vertex in self.adj_list[v] and vertex in visited:
                cols_temp.add(vertex.col)
        for c in cols:
            if c not in cols_temp:
                return c
        cols.append(cols[-1] + 1)
        return cols[-1]

    def __dfs_color(self, vertex: Vertex) -> None:
        if not self.is_empty() and vertex in self.adj_list:
            cols = [1]
            visited = set()

            def dfs_color_in(v: Vertex) -> None:
                if v not in visited:
                    v.col = self.__check_color_availability(v, visited, cols)
                    visited.add(v)
                    for neigh in self.list_vertex:
                        if neigh not in visited and neigh in self.adj_list[v]:
                            dfs_color_in(neigh)

            dfs_color_in(vertex)

    def __bfs_color(self, vertex: Vertex) -> None:
        if not self.is_empty() and vertex in self.adj_list:
            cols = [1]
            visited = set()
            vertex.col = self.__check_color_availability(vertex, visited, cols)
            visited.add(vertex)
            queue = [vertex]
            while queue:
                popped = queue.pop(0)
                for neigh in self.list_vertex:
                    if neigh not in visited and neigh in self.adj_list[popped]:
                        neigh.col = self.__check_color_availability(neigh, visited, cols)
                        visited.add(neigh)
                        queue.append(neigh)

    def color(self, vertex: Vertex, t: str) -> None:
        if t == 'dfs':
            self.__dfs_color(vertex)
        elif t == 'bfs':
            self.__bfs_color(vertex)


def color_graph(G: Graph, vertex: Vertex, t: str) -> None:
    G.color(vertex, t)


def printGraph(g):
    n = g.size()
    print("------GRAPH------ ", n)
    for i in range(n):
        v = g.list_vertex[i]
        print(v, end=" -> ")
        for ver, w in g.adj_list[v].items():
            print(ver, w, end="; ")
        print()
    print("-------------------")


def main():
    graph = [('A', 'B', 4), ('A', 'C', 1), ('A', 'D', 4),
            ('B', 'E', 9), ('B', 'F', 9), ('B', 'G', 7), ('B', 'C', 5),
            ('C', 'G', 9), ('C', 'D', 3),
            ('D', 'G', 10), ('D', 'J', 18),
            ('E', 'I', 6), ('E', 'H', 4), ('E', 'F', 2),
            ('F', 'H', 2), ('F', 'G', 8),
            ('G', 'H', 9), ('G', 'J', 8),
            ('H', 'I', 3), ('H', 'J', 9),
            ('I', 'J', 9)
            ]
    G = GraphList()
    for t in graph:
        G.insert_edge(Vertex(t[0]), Vertex(t[1]), t[2])
    G_prim, length = G.prim(G.list_vertex[0])
    printGraph(G_prim)


if __name__ == '__main__':
    main()