1579. Remove Max Number of Edges to Keep Graph Fully Traversable - LeetCode

vertexT = int


class Uf1:
    def __init__(self, root: list[vertexT], groups: int):
        self.root = root
        self.groups = groups

    def find(self, x: vertexT):
        if self.root[x] != x:
            self.root[x] = self.find(self.root[x])
        return self.root[x]

    def union(self, x: vertexT, y: vertexT):
        x_root, y_root = self.find(x), self.find(y)
        result = int(x_root != y_root)
        if result:
            self.root[x_root] = y_root
            self.groups -= 1
        return result


class Solution1:
    def maxNumEdgesToRemove(self, n: int, edges: list[list[int]]) -> int:
        result = 0
        ufs = [Uf1(list(range(n)), n), Uf1(list(range(n)), n)]
        for t, u, v in sorted(edges, reverse=True):
            r = 1
            for i in range(len(ufs) - 1, -1, -1):
                if t > i:
                    r &= ufs[i].union(u - 1, v - 1)
                    t -= i + 1
            result += r ^ 1  # 1, 0 -> 0, 1
        return result if ufs[0].groups == ufs[1].groups == 1 else -1


class Uf:
    def __init__(self, root: list[vertexT]):
        self.root = root

    def find(self, x: vertexT):
        if self.root[x] != x:
            self.root[x] = self.find(self.root[x])
        return self.root[x]

    def union(self, x: vertexT, y: vertexT):
        x_root, y_root = self.find(x), self.find(y)
        result = int(x_root != y_root)
        if result:
            self.root[x_root] = y_root
        return result


class Solution:
    def maxNumEdgesToRemove(self, n: int, edges: list[list[int]]) -> int:
        result = 0
        ufs = [Uf(list(range(n))), Uf(list(range(n)))]
        groups = [n, n]
        for t, u, v in sorted(edges, reverse=True):
            r = 1
            for i in range(len(groups) - 1, -1, -1):
                if t > i:
                    uni = ufs[i].union(u - 1, v - 1)
                    if uni:
                        groups[i] -= 1
                    r &= uni
                    t -= i + 1
            result += r ^ 1  # 1, 0 -> 0, 1
        return result if all(x == 1 for x in groups) else -1