python实现的TopologySort

from collections import defaultdict, deque

class Graph:
    def __init__(self):
        self.graph = defaultdict(list)
        self.in_degree_dict = defaultdict(int)
        self.result = []

    def add_edge(self, u, v):
        self.graph[u].append(v)                             # 构造拓扑图
        if(self.in_degree_dict.get(u) == None):             # 统计每个节点的入度
            self.in_degree_dict[u] = 0
        if(self.in_degree_dict.get(v) == None):
            self.in_degree_dict[v] = 1
        else:
            val = self.in_degree_dict.get(v) + 1
            self.in_degree_dict[v] = val

    def topological_sort(self):
        # print('-----graphMsg------>', type(self.graph))
        # print('-----graphMsg------>', self.in_degree_dict)
        queue = deque()         # 初始化一个队列用于存放入度为0的节点
        for key in self.in_degree_dict:
            if self.in_degree_dict[key] == 0:
                queue.append(key)

        while queue:
            node = queue.popleft()
            self.result.append(node)

            # 访问每一个入度=0节点的相邻node，每次被访问到的节点，就对其入度-1;
            # 入度为0的节点加入queue当中，等待继续访问它们的相邻node
            for neighbor in self.graph[node]:
                val = self.in_degree_dict.get(neighbor) - 1
                self.in_degree_dict[neighbor] = val
                if self.in_degree_dict[neighbor] == 0:
                    queue.append(neighbor)

        if len(self.result) == len(self.in_degree_dict):         # 如果结果中的节点数量等于图中的节点数量，则表示拓扑排序成功
            return "拓扑排序结果:", self.result
        else:
            return "图中存在环，无法进行拓扑排序。", []

    def get_work_pair(self):        # 根据拓扑排序的结果生成工作序列
        pair_list = []
        if len(self.result) == len(self.in_degree_dict) and len(self.result) != 0:
            for node in self.result:
                sub_list = self.graph.get(node)
                for next in sub_list:
                    # print(node, next)
                    pair_list.append((node, next))
        return pair_list


def main():
    g = Graph()
    g.add_edge(1, 4)
    g.add_edge(1, 2)
    g.add_edge(2, 4)
    g.add_edge(4, 3)
    g.add_edge(4, 5)
    g.add_edge(3, 5)
    g.add_edge(2, 5)
    g.add_edge(5, 6)
    g.add_edge(3, 6)
    g.add_edge(6, 7)
    # g.add_edge(7, 1)

    topo_list = g.topological_sort()
    print(topo_list)
    pair_list = g.get_work_pair()
    print(pair_list)

if __name__ == "__main__":
    main()