Leetcode 1701 Solution

This article provides solution to leetcode question 1701 (remove-max-number-of-edges-to-keep-graph-fully-traversable)
Categories:
Tags:
Access this page by simply typing in "lcs 1701" in your browser address bar if you have bunnylol configured.

https://leetcode.com/problems/remove-max-number-of-edges-to-keep-graph-fully-traversable

Solution

class FindUnionSet:
    def __init__(self, n):
        self.n = n
        self.p = list(range(n))

    def copy(self):
        fmset = FindUnionSet(self.n)
        fmset.p = list(self.p)
        return fmset

    def union(self, i, j):
        ri = self.find(i)
        rj = self.find(j)

        self.p[ri] = rj

    def find(self, i):
        if self.p[i] == i:
            return self.p[i]

        self.p[i] = self.find(self.p[i])
        return self.p[i]


class Solution:
    def maxNumEdgesToRemove(self, n: int, edges: List[List[int]]) -> int:
        base_fmset = FindUnionSet(n + 1)

        ans = 0
        for t, src, dst in edges:
            if t == 3:
                if base_fmset.find(src) == base_fmset.find(dst):
                    ans += 1
                else:
                    base_fmset.union(src, dst)

        fmset1 = base_fmset.copy()
        for t, src, dst in edges:
            if t == 1:
                if fmset1.find(src) == fmset1.find(dst):
                    ans += 1
                else:
                    fmset1.union(src, dst)

        root_set1 = {fmset1.find(i) for i in range(1, n + 1)}
        if len(root_set1) != 1:
            return -1

        fmset2 = base_fmset.copy()
        for t, src, dst in edges:
            if t == 2:
                if fmset2.find(src) == fmset2.find(dst):
                    ans += 1
                else:
                    fmset2.union(src, dst)


        root_set1 = {fmset2.find(i) for i in range(1, n + 1)}
        if len(root_set1) != 1:
            return -1

        return ans