1489. 找到最小生成树里的关键边和伪关键边

题目描述：给你一个 n 个点的带权无向连通图，节点编号为 0 到 n-1 ，同时还有一个数组 edges ，其中 edges[i] = [fromi, toi, weighti] 表示在 fromi 和 toi 节点之间有一条带权无向边。最小生成树 (MST) 是给定图中边的一个子集，它连接了所有节点且没有环，而且这些边的权值和最小。
请你找到给定图中最小生成树的所有关键边和伪关键边。如果从图中删去某条边，会导致最小生成树的权值和增加，那么我们就说它是一条关键边。伪关键边则是可能会出现在某些最小生成树中但不会出现在所有最小生成树中的边。
请注意，你可以分别以任意顺序返回关键边的下标和伪关键边的下标。
解题思路：Kruskal算法+并查集+Tarjan算法

class UnionFind():
    def __init__(self, n):
        self.rank = [1] * n
        self.f = list(range(n))
        
    def find(self, x):
        root = x
        while(root != self.f[root]):
            root = self.f[root]
        while(x != root):
            origin_r = self.f[x]
            self.f[x] = root
            x = origin_r
            
        return root
    
    def union(self, x, y):
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:
            root_x, root_y = root_y, root_x
        self.rank[root_x] += self.rank[root_y]
        self.f[root_y] = root_x
        return True

class Tarjan():
    
    def __init__(self, n, vids, eids):
        self.n = n
        self.vids = vids
        self.eids = eids
        self.low = [-1] * n
        self.dfn = [-1] * n
        self.ans = []
        self.clock = -1
        
    def getCuttingEdge(self):
        for i in range(self.n):
            if self.dfn[i] == -1:
                self.pGetCuttingEdge(i, -1)
        return self.ans
    
    
    def pGetCuttingEdge(self, u, parentEdgeid):
        self.clock += 1
        self.dfn[u] = self.low[u] = self.clock
        for v, iden in zip(self.vids[u], self.eids[u]):
            if self.dfn[v] == -1:
                self.pGetCuttingEdge(v, iden)
                self.low[u] = min(self.low[u], self.low[v])
                if self.low[v] > self.dfn[u]:
                    self.ans.append(iden)
            elif iden != parentEdgeid:
                self.low[u] = min(self.low[u], self.dfn[v])
class Solution:
    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        from collections import defaultdict
        uf = UnionFind(n)
        m = len(edges)
        for i in range(m):
            edges[i].append(i)
        edges.sort(key=lambda x: x[2])
        
        label = [0] * m #表示是否为关键边:1代表关键边，0表示伪关键边，-1表示可以舍弃的边
        i = 0
        ans0 = []
        while(i < m):
            w = edges[i][2]
            j = i
            while(j < m and edges[j][2] == w):
                j += 1
            gn = 0
            g = {}
            for k in range(i, j):
                x, y = edges[k][0], edges[k][1]
                rootx, rooty = uf.find(x), uf.find(y)
                if rootx == rooty:
                    label[edges[k][3]] = -1
                else:
                    if rootx not in g:
                        g[rootx] = gn
                        gn += 1
                    if rooty not in g:
                        g[rooty] = gn
                        gn += 1
            eid = defaultdict(list)
            vid = defaultdict(list)
            for k in range(i, j):
                x, y = edges[k][0], edges[k][1]
                rootx, rooty = uf.find(x), uf.find(y)
                if rootx != rooty:
                    xid, yid = g[rootx], g[rooty]
                    eid[xid].append(edges[k][3])
                    eid[yid].append(edges[k][3])
                    vid[xid].append(yid)
                    vid[yid].append(xid)
            
            edge = Tarjan(gn, vid, eid).getCuttingEdge()      
            ans0.extend(edge)
            for e in edge:
                label[e] = 1
            for k in range(i, j):
                x, y = edges[k][0], edges[k][1]
                uf.union(x, y)
            i = j
        ans1 = [i for i in range(m) if label[i] == 0]
        return [ans0, ans1]