LC 2076 [H] Process Restricted Friend Requests - ALawliet/algorithms GitHub Wiki

class DSU:
    def __init__(self, N):
        self.parent = list(range(N))

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

    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        self.parent[x_root] = y_root # x becomes child of y

class Solution:
    def friendRequests(self, n, restrictions, requests):
            dsu = DSU(n)
            res = []
            
            for i, j in requests:
                success = True
                pi, pj = dsu.find(i), dsu.find(j)
                
                if pi != pj:
                    for x, y in restrictions:
                        px, py = dsu.find(x), dsu.find(y)
                        
                        if (pi, pj) == (px, py) or (pj, pi) == (px, py):
                            success = False
                            break
                            
                if success:
                    dsu.union(pj, pi)
                    
                res.append(success)
                
            return res