Leetcode: 1579. Remove Max Number of Edges to Keep Graph Fully Traversable

Blackbox: This is a variation of a classic problem to find Connected Components in a Graph.
Can we derive an invariant based on the smallest possible examples? Be \(G\) a graph and \(T\) a tree of \(G\). Be \(u\) and \(v\) in \(T\) where they need two edges of type 1 and 2. If it is possible to swap these edges with one of type 3, then \(|E(G)|-|E(T)|\) is not the solution since we can do better. Therefore, we can first add all edges of type 3 and then handle the ones with type 1 and type 2. Time and space complexity are \(O(n)\).

from typing import List


class Solution:
    def maxNumEdgesToRemove(self, n: int, edges: List[List[int]]) -> int:
        p = {}

        def find(u):
            if u not in p:
                p[u] = u
            if p[u] != u:
                p[u] = find(p[u])
            return p[u]

        def union(u, v):
            p[find(u)] = find(v)

        added = 0
        for t, u, v in sorted(edges, key=lambda e: e[0], reverse=True):
            if t == 3:
                if find(u) != find(v):
                    union(u, v)
                    union(~u, ~v)
                    added += 1
            elif t == 2:
                if find(~u) != find(~v):
                    union(~u, ~v)
                    added += 1
            else:
                if find(u) != find(v):
                    union(u, v)
                    added += 1

        valid = (
            len({find(u) for u in range(1, n + 1)})
            == len({find(~u) for u in range(1, n + 1)})
            == 1
        )
        return -1 if not valid else len(edges) - added


assert (
    Solution().maxNumEdgesToRemove(
        4, [[3, 1, 2], [3, 2, 3], [1, 1, 3], [1, 2, 4], [1, 1, 2], [2, 3, 4]]
    )
    == 2
)
assert (
    Solution().maxNumEdgesToRemove(4, [[3, 1, 2], [3, 2, 3], [1, 1, 4], [2, 1, 4]]) == 0
)
assert Solution().maxNumEdgesToRemove(4, [[3, 2, 3], [1, 1, 2], [2, 3, 4]]) == -1