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

from typing import List
# 这道题本质上还是利用并查集的知识，首先计算最小的权重和，去除掉无效的边，
# 然后遍历所有的边，然后计算最小的权重和，同时判断是否可以连通。

# 并查集的模板。
class DSU:
    def __init__(self,n):
        # 实例化一个列表。用来存放所有的节点对应的父节点。
        self.father = list(range(n))
    def find(self,x):
        # 寻找父节点，判断父节点是否是自己。
        if x != self.father[x]:
            # 寻找到最上层的父亲节点。
            self.father[x] = self.find(self.father[x])
        # 返回。
        return self.father[x]
    def union(self,x,y):
        # 联合，连接节点。
        if not self.is_connected(x,y):
            # 寻找到最上层的x的父亲节点，让他等于y的父亲节点。
            self.father[self.find(x)] = self.find(y)
    # 判断两个节点是否相连。
    def is_connected(self,x,y):
        return self.find(x) == self.find(y)


class Solution:
    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        # 在原有的列表中添加位置索引。
        sorted_edges = [[index] + i for index,i in enumerate(edges)]
        # 根据边的权重对边进行排序。
        sorted_edges = sorted(sorted_edges,key = lambda x : x[-1])

        # 计算最小生成树的权重和。
        total = 0
        dsu = DSU(n)
        for index2,x,y,w in sorted_edges:
            if not dsu.is_connected(x,y):
                dsu.union(x,y)
                total += w
        key_edge = [] # 关键边列表
        no_key_edge = [] # 伪关键边列表。
        for i,edge in enumerate(sorted_edges):
            (j,x,y,w) = edge
            # 去掉当前边，形成新的列表。
            cur_edges = sorted_edges[:i] + sorted_edges[i+1:]
            
            # 这里我们先把当前边相连。然后在进行计算权重和，
            # 判断权重和和最小权重和是否相同。
            # 如果不相同的话，那么这条边就是无效的边。所以就跳出这一次循环。
            total1 = 0
            dsu1 = DSU(n)
            dsu1.union(x,y)
            total1 += w
            for index2, cur_x, cur_y, cur_w in cur_edges:
                if not dsu1.is_connected(cur_x, cur_y):
                    dsu1.union(cur_x, cur_y)
                    total1 += cur_w
            # 如果不相等说明是无效边。
            if total1 != total:
                continue
                
            # 下边是判断是否为关键边。
            # 下边我们不将当前边对应的两个节点相连。
            # 然后进行计算最小权重和，如果当前边是关键边的话，
            # 那么我们是不能保证能够让所有的节点相连的。
            # 所以计算的权重和total2 和 total1是不相同的。
            total2 = 0
            dsu2 = DSU(n)
            for index2, cur_x, cur_y, cur_w in cur_edges:
                if not dsu2.is_connected(cur_x, cur_y):
                    dsu2.union(cur_x, cur_y)
                    total2 += cur_w
            # 判断是否为关键边。
            if total1 != total2:
                key_edge.append(edge[0])
            else:
                no_key_edge.append(edge[0])
        return [key_edge,no_key_edge]

A = Solution()
print(A.findCriticalAndPseudoCriticalEdges(n = 5, edges = [[0,1,1],[1,2,1],[2,3,2],[0,3,2],[0,4,3],[3,4,3],[1,4,6]]))