Python中使用朱刘算法解决最小树形图问题

import sys

class Edge:
    def __init__(self, u, v, w):
        self.u = u
        self.v = v
        self.w = w

class Graph:
    def __init__(self, n, edges):
        self.n = n
        self.edges = edges

    def add_edge(self, u, v, w):
        self.edges.append(Edge(u, v, w))

    def find_cycle(self):
        self.id = [-1] * self.n
        self.vis = [-1] * self.n
        for i in range(self.n):
            if self.vis[i] == -1 and self.dfs(i):
                return True
        return False

    def dfs(self, u):
        self.vis[u] = 1
        for e in self.edges:
            if e.u != u:
                continue
            v = e.v
            w = e.w
            if self.vis[v] == -1:
                if self.dfs(v):
                    return True
            elif self.id[v] == -1:
                self.id[v] = u
                self.val[v] = w
            elif self.val[v] > w:
                self.id[v] = u
                self.val[v] = w
        self.vis[u] = 2
        return False

    def min_tree(self, root):
        res = 0
        while True:
            self.id = [-1] * self.n
            self.val = [sys.maxsize] * self.n
            if self.find_cycle():
                for i in range(self.n):
                    if self.vis[i] == 2:
                        continue
                    if self.id[i] == -1:
                        self.id[i] = root
                    res += self.val[i]
                self.compress()
            else:
                break
        return res

    def compress(self):
        nc = 0
        nedges = []
        for i in range(self.n):
            self.vis[i] = -1
        for i in range(self.n):
            if self.vis[i] >= 0:
                continue
            q = [i]
            while q:
                u = q.pop(0)
                if self.vis[u] >= 0:
                    continue
                self.vis[u] = nc
                for e in self.edges:
                    if e.u != u:
                        continue
                    v = e.v
                    if self.id[v] == u:
                        continue
                    if self.vis[v] == -1:
                        nedges.append(Edge(nc, self.vis[v], e.w - self.val[v]))
                        q.append(v)
                    else:
                        nedges.append(Edge(nc, self.vis[v], e.w))
            nc += 1
        self.edges = nedges
        self.n = nc

if __name__ == '__main__':
    n = 6
    edges = []
    edges.append(Edge(0, 1, 1))
    edges.append(Edge(1, 2, 2))
    edges.append(Edge(2, 3, 3))
    edges.append(Edge(3, 1, -6))
    edges.append(Edge(4, 5, 1))
    edges.append(Edge(5, 4, 2))
    edges.append(Edge(3, 5, 3))
    g = Graph(n, edges)
    print(g.min_tree(0))  # 输出为-5，表示最小树形图的边权之和为-5