Redundant Connection II

There are two cases which make tree invalid:

• more than one parents
• has circle

The algorithm contains two steps:

1. Check whether there is a node has more than one parents:
   If so, mark these two edges as candidateA and candidateB, and make the second edge invalid.

2. Perform regular union and find:
   If tree is valid

   • simply return candidateB (since we invalidate candidateB and the remaining tree is valid)
     If tree contains a circle
   • if candidateA is not null, return candidateA
   • if candidateA is null, return this edge

class Solution(object):
    def findRedundantDirectedConnection(self, edges):
        """
        :type edges: List[List[int]]
        :rtype: List[int]
        """
        graph = DSU(len(edges)+1)
        self.parent = [0] * (len(edges)+1)
        candidateA = None
        candidateB = None

        # find whether a node has more than one parents
        for edge in edges:
            if self.parent[edge[1]] == 0:
                self.parent[edge[1]] = edge[0]
            else:
                candidateA = [self.parent[edge[1]], edge[1]][:]
                candidateB = edge[:]
                edge[1] = -1 # set second edge invalid
                #print candidateA
                #print candidateB

        # find circle
        for edge in edges:
            v1 = edge[0]
            v2 = edge[1]

            if v2 == -1:
                continue

            if graph.find(v1) == graph.find(v2):
                if candidateA:
                    return candidateA
                return edge
            graph.union(v1, v2)

        return candidateB


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

    def add(self, x):
        self.par[x] = x

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

    def union(self, x, y):
        self.par[self.find(x)] = self.find(y)