'''
由于输入数据是有向边, 所以使用并查集更合适
BFS还需要先把有向边转换成无向边

431： 强连通分量，432：弱连通分量。
图的节点不一定是连通的。
可以用BFS, 但是跟431不相同的点在于，
431【这个点没有被访问过，说明这个点的连通分量也没到】是可以成立的，但是这个题不行，比如从2可以到4，但是4不一定能到2，
所以先需要改造一下图，把所有的有向边都加成无向边，然后就可以用431的做法了。
'''

class DirectedGraphNode:
     def __init__(self, x):
         self.label = x
         self.neighbors = []

class UnionFind:
    def __init__(self, nodes):
        self.father = {}
        for node in nodes:
            self.father[node.label] = node.label

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

    def union(self, a, b):
        root_a = self.find(a)
        root_b = self.find(b)
        if root_a != root_b:
            self.father[root_a] = root_b


class Solution:
    """
    @param nodes: a array of Directed graph node
    @return: a connected set of a directed graph
    """

    def connectedSet2(self, nodes):
        results = []
        # self.father[node.label] = node.label
        uf = UnionFind(nodes)

        # 建立并查集
        for node in nodes:
            for n in node.neighbors:
                uf.union(n.label, node.label)

        g = {}
        cnt = 0

        # 存放root， 用SET避免重复
        myset = set()

        for node1 in nodes:
            # find每个结点的ROOT,并且放进集合里
            node1_root = uf.find(node1.label)
            if node1_root not in myset:
                myset.add(node1_root)
                # 找以x为根结点的所有结点，放入LS中
                ls = []
                for node2 in nodes:
                    node2_root = uf.find(node2.label)
                    if node2_root == node1_root:
                        ls.append(node2.label)
                results.append(ls)

        # for node in nodes:
        #     x = uf.find(node.label)
        #     if not x in g:
        #         cnt += 1
        #         g[x] = cnt
        #
        #     if len(results) < cnt:
        #         results.append([])
        #
        #     results[g[x] - 1].append(node.label)
        return results


if __name__ == '__main__':
    sol = Solution()
    node1 = DirectedGraphNode(1)
    node2 = DirectedGraphNode(2)
    node3 = DirectedGraphNode(3)
    node4 = DirectedGraphNode(4)
    node5 = DirectedGraphNode(5)
    node6 = DirectedGraphNode(6)

    node1.neighbors.append(node2)
    node1.neighbors.append(node4)

    node2.neighbors.append(node4)

    node3.neighbors.append(node5)

    node6.neighbors.append(node5)

    # node4.neighbors.append(node1)
    # node4.neighbors.append(node2)
    #
    # node5.neighbors.append(node3)

    nodes = [node1, node2, node3, node4, node5, node6]
    res = sol.connectedSet2(nodes)
    print(res)