hello-algo/codes/python/chapter_graph/graph_adjacency_list.py

"""
File: graph_adjacency_list.py
Created Time: 2023-02-23
Author: Krahets (krahets@163.com)
"""

import sys, os.path as osp
sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
from modules import *

class GraphAdjList:
    """ 基于邻接表实现的无向图类 """
    def __init__(self, edges: List[List[Vertex]]) -> None:
        """ 构造方法 """
        # 邻接表，key: 顶点，value：该顶点的所有邻接顶点
        self.adj_list: Dict = {}
        # 添加所有顶点和边
        for edge in edges:
            self.add_vertex(edge[0])
            self.add_vertex(edge[1])
            self.add_edge(edge[0], edge[1])

    def size(self) -> int:
        """ 获取顶点数量 """
        return len(self.adj_list)

    def add_edge(self, vet1: Vertex, vet2: Vertex) -> None:
        """ 添加边 """
        if vet1 not in self.adj_list or vet2 not in self.adj_list or vet1 == vet2:
            raise ValueError
        # 添加边 vet1 - vet2
        self.adj_list[vet1].append(vet2)
        self.adj_list[vet2].append(vet1)

    def remove_edge(self, vet1: Vertex, vet2: Vertex) -> None:
        """ 删除边 """
        if vet1 not in self.adj_list or vet2 not in self.adj_list or vet1 == vet2:
            raise ValueError
        # 删除边 vet1 - vet2
        self.adj_list[vet1].remove(vet2)
        self.adj_list[vet2].remove(vet1)

    def add_vertex(self, vet: Vertex) -> None:
        """ 添加顶点 """
        if vet in self.adj_list:
            return
        # 在邻接表中添加一个新链表
        self.adj_list[vet] = []

    def remove_vertex(self, vet: Vertex) -> None:
        """ 删除顶点 """
        if vet not in self.adj_list:
            raise ValueError
        # 在邻接表中删除顶点 vet 对应的链表
        self.adj_list.pop(vet)
        # 遍历其它顶点的链表，删除所有包含 vet 的边
        for vertex in self.adj_list:
            if vet in self.adj_list[vertex]:
                self.adj_list[vertex].remove(vet)

    def print(self) -> None:
        """ 打印邻接表 """
        print("邻接表 =")
        for vertex in self.adj_list:
            tmp = [v.val for v in self.adj_list[vertex]]
            print(f"{vertex.val}: {tmp},")


""" Driver Code """
if __name__ == "__main__":
    """ 初始化无向图 """
    v = vals_to_vets([1, 3, 2, 5, 4])
    edges = [[v[0], v[1]], [v[0], v[3]], [v[1], v[2]],
             [v[2], v[3]], [v[2], v[4]], [v[3], v[4]]]
    graph = GraphAdjList(edges)
    print("\n初始化后，图为")
    graph.print()

    """ 添加边 """
    # 顶点 1, 2 即 v[0], v[2]
    graph.add_edge(v[0], v[2])
    print("\n添加边 1-2 后，图为")
    graph.print()

    """ 删除边 """
    # 顶点 1, 3 即 v[0], v[1]
    graph.remove_edge(v[0], v[1])
    print("\n删除边 1-3 后，图为")
    graph.print()

    """ 添加顶点 """
    v5 = Vertex(6)
    graph.add_vertex(v5)
    print("\n添加顶点 6 后，图为")
    graph.print()

    """ 删除顶点 """
    # 顶点 3 即 v[1]
    graph.remove_vertex(v[1])
    print("\n删除顶点 3 后，图为")
    graph.print()