leetcode 1579. 保证图可完全遍历 python 并查集

Alice 和 Bob 共有一个无向图，其中包含 n 个节点和 3 种类型的边：

类型 1：只能由 Alice 遍历。
类型 2：只能由 Bob 遍历。
类型 3：Alice 和 Bob 都可以遍历。
给你一个数组 edges ，其中 edges[i] = [typei, ui, vi] 表示节点 ui 和 vi 之间存在类型为 typei 的双向边。请你在保证图仍能够被 Alice和 Bob 完全遍历的前提下，找出可以删除的最大边数。如果从任何节点开始，Alice 和 Bob 都可以到达所有其他节点，则认为图是可以完全遍历的。

返回可以删除的最大边数，如果 Alice 和 Bob 无法完全遍历图，则返回 -1 。

解法：先将类型3的弄到并查集中，然后将类型1、2的弄进去，如果已经连通，则说明这条是多余的。

class UnionFindSet():
    def __init__(self,n):
        self.setSize = n  #不连通区域
        self.father = {}

    def add(self,x):  #添加根节点
        if x not in self.father:
            self.father[x] = -1
            #self.setSize += 1

    def find(self,x):
        root = x
        while self.father[root] != -1:  #找根节点
            root = self.father[root]

        while (x != root): #路径压缩
            o = self.father[x]  #找x的父节点
            self.father[x] = root  #把x的父节点设置成刚才找到的根
            x = o  #往上一层

        return root

    def merge(self,x, y):
        self.add(x)
        self.add(y)
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x == root_y:
            return False
        else:
            self.father[root_x] = root_y  #合并
            self.setSize -= 1
            return True
            
    def is_Union(self, x, y):
        return self.find(x) == self.find(y)


class Solution:
    def maxNumEdgesToRemove(self, n: int, edges: List[List[int]]) -> int:
        uf_a, uf_b = UnionFindSet(n), UnionFindSet(n)
        ans = 0

        for t, i, j in edges:
            if t == 3:
                if not uf_a.merge(i, j):
                    ans += 1
                else:
                    uf_b.merge(i, j)
        for t, i, j in edges:
            if t == 1:
                if not uf_a.merge(i, j):
                    ans += 1
            elif t == 2:
                if not uf_b.merge(i, j):
                    ans += 1
        if uf_a.setSize != 1 or uf_b.setSize != 1:
            return - 1
        else:
            return ans