graphClass.py

import matplotlib.pyplot as plt
import networkx as nx

class Graph:
    def __init__(self, directed: bool = False, weighted: bool = False) -> None:
        self.directed = directed
        self.weighted = weighted
        self.adjacency_dict = {}

    def __str__(self) -> str:
        return str(self.adjacency_dict)

    def add_vertex(self, key: str) -> None:
        if not key:
            raise ValueError("Key is empty")
        if key in self.adjacency_dict:
            raise ValueError("Key already exists")

        self.adjacency_dict[key] = {} if self.weighted else set()

    def add_edge(self, key: str, value: str, weight: float | None = None) -> None:
        if not key or not value:
            raise ValueError("Key/value is empty")
        if key not in self.adjacency_dict:
            raise ValueError(f"Key does not exist: {key}")
        if value not in self.adjacency_dict:
            raise ValueError(f"Value does not exist: {value}")

        if self.weighted:
            if weight is None:
                raise ValueError("Weight is empty")
            if not isinstance(weight, (int, float)):
                raise ValueError("Weight must be a number")
            self.adjacency_dict[key][value] = float(weight)
            if not self.directed:
                self.adjacency_dict[value][key] = float(weight)
        else:
            self.adjacency_dict[key].add(value)
            if not self.directed:
                self.adjacency_dict[value].add(key)

    def neighbors(self, key: str) -> list:
        if key not in self.adjacency_dict:
            raise ValueError(f"Vertex does not exist: {key}")
        return self.adjacency_dict[key]

    def deg(self, key: str) -> int:
        if key not in self.adjacency_dict:
            raise ValueError(f"Vertex does not exist: {key}")
        return len(self.adjacency_dict[key])
    
    def visualize_graph(self) -> None:
        G = nx.DiGraph() if self.directed else nx.Graph()

        for vertex in self.adjacency_dict.keys():
            G.add_node(vertex)

        for vertex, edges in self.adjacency_dict.items():
            if self.weighted:
                for edge, weight in edges.items():
                    G.add_edge(vertex, edge, weight=weight)
            else:
                for edge in edges:
                    G.add_edge(vertex, edge)

        pos = nx.spring_layout(G)
        nx.draw(
            G, pos,
            with_labels=True,
            node_color="lightblue",
            node_size=700,
            font_size=16,
            font_color="black"
        )
        if self.weighted:
            edge_labels = nx.get_edge_attributes(G, "weight")
            nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
        plt.show()

class NormalGraph(Graph):
    def __init__(self) -> None:
        super().__init__(directed=False, weighted=False)

class DirectedGraph(Graph):
    def __init__(self) -> None:
        super().__init__(directed=True, weighted=False)

class WeightedGraph(Graph):
    def __init__(self) -> None:
        super().__init__(directed=False, weighted=True)

class DirectedWeightedGraph(Graph):
    def __init__(self) -> None:
        super().__init__(directed=True, weighted=True)