1489.-find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree.md

1489. Find Critical and Pseudo-Critical Edges in Minimum Spanning Tree

class Solution:
    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        m = len(edges)

        edges = [edges[i] + [i] for i in range(m)]
        edges = sorted(edges, key = lambda x : x[2])
        print(edges)
        def inital():
            f = [i for i in range(n)]
            size = [1 for i in range(n)]
            return f, size 
        f , size = inital()
        def find(x):
            if f[x] != x :
                f[x] = find(f[x])

            return f[x]

        def union(x , y):
            fx = find(x)
            fy = find(y)
            if fx == fy :
                return False
            if size[fy] > size[fx]:
                size[fy] += size[fx] 
                f[fx] = fy
            else:
                size[fx] += size[fy] 
                f[fy] = fx
            return True


        
        min_value = 0
        for i in range(m):
            ## 
            if union(edges[i][0], edges[i][1]):
                min_value += edges[i][2]
        
        ans = [[],[]]
        for i in range(m):
            value = 0
            f , size = inital()
            for j in range(m):
                if i != j and union(edges[j][0], edges[j][1]):
                    value += edges[j][2]

            if max(size) < n  or (max(size) == n  and value > min_value):
                ans[0].append(edges[i][3])
                continue

            f , size = inital()
            union(edges[i][0], edges[i][1])
            v = edges[i][2]
            for j in range(m):
                if i != j and union(edges[j][0], edges[j][1]):
                    v += edges[j][2]
            if value == v:
                ans[1].append(edges[i][3])
        return ans