def undirected_edge_index(
    edge_index
):
    if isinstance(edge_index,list):
        edge_index = np.array(edge_index)

    # 接続元と接続先を入れ替えたものを追記する
    edge_index = np.concatenate([
        edge_index,
        edge_index[:,[1,0]]
    ])
    
    # 重複を除く
    edge_index = np.unique(edge_index, axis=0)
    
    return edge_index

edge_index_filtered = undirected_edge_index(edge_index_filtered)

adj_matrix_undirected = adj_matrix + adj_matrix.T
adj_matrix_undirected[adj_matrix_undirected > 0] = 1

def find_connected_cluster(start_node, edge_index):
    """
    1つのスタートノードから辿り着くことができる部分グラフのエッジを取り出す機能
    Return:
      サブグラフのノード番号リスト
    """
    
    def get_connected_edges(start_node, edge_index):
        connected_edges = [ edge for edge in edge_index if start_node in edge]
        return connected_edges
    
    # 探索済みか記録する用　ノード数の長さのリストにbool値を記録
    num_nodes = np.unique(list(edge_index)).max()
    visited_nodes = np.zeros(num_nodes, dtype=bool)
    # start nodeの箇所をTrueにしておく
    visited_nodes[start_node] = True

    # 最終出力のリスト。サブグラフ内の全エッジを記録
    cluster = []
    connected_nodes = {start_node}

    ## 接続先ノードが無くなるまで繰り返す
    while connected_nodes:
        new_edges = []
        new_nodes = set()

        for node in sorted(connected_nodes):
            # 1つのノードに繋がるエッジを取得
            edges = get_connected_edges(node, edge_index)
            # 各エッジが探索済みノードのものかを判定。未探索であれば記録
            for edge in edges:
                if not (visited_nodes[edge[0]] and visited_nodes[edge[1]]):
                    new_edges.append(edge)
                    new_nodes.update(edge)
        
        # 未探索のエッジがなければwhileを終了
        if not new_edges:
            break

        # clusterにエッジを記録
        cluster.extend(new_edges)
        # 初探索ノードを更新して次のloopに渡す
        connected_nodes = new_nodes - set(np.where(visited_nodes)[0]) # 探索済みは除く
        visited_nodes[list(connected_nodes)] = True # このloopで探索したノードは探索済みとする

    return cluster

cluster_edges = find_connected_cluster(start_node=0, edge_index=edge_index_filtered)

cluster_nodes = np.unique(cluster_edges)

def find_clusters(
    edge_index
):
    clusters = []
    nodes = np.unique(edge_index)
    visited = set() # 探索済みノードかどうかの記録用

    for n in nodes:
        # 探索済みならskip
        if n not in visited:
            # ノードnに接続されたエッジを見つける
            edges = find_connected_cluster(n, edge_index)

            if edges:
                # クラスターに属するノードを見つける
                cluster_nodes = list(np.sort(np.unique(edges)))
                clusters.append(cluster_nodes)

                # クラスターのノードをvisitedに追加
                visited.update(cluster_nodes)
            else:
                # 孤立ノードとして個別に追加
                clusters.append([n])
                visited.add(n)

    return clusters

clusters = find_clusters(edge_index_filtered)

class UnionFind:
    def __init__(self, n_node):
        self.n_node = n_node
        # 初期クラスター番号（nodeの数だけ用意し、ノード番号=クラスター番号の状態）
        self.parent = list(range(n_node))
        # 各ノードの最長経路長を保存する用
        self.rank = [1] * n_node
    
    def find(self, u):
        """
        ノード番号uのクラスター番号を見つける機能
        初期状態であればノード番号uのクラスター番号はu
        クラスター統合が進むと、ノード番号uのクラスター番号が変わっている。
        ノード番号uのクラスター番号を再帰的に探す。
        """
        # ノード番号uのクラスター番号がuでなければ
        if self.parent[u] != u:
            # ノード番号 == クラスター番号の箇所を再帰的に探して、そのクラスター番号を上書き
            self.parent[u] = self.find(self.parent[u])
        
        # ノード番号uのクラスター番号を返す
        return self.parent[u]
    
    def union(self, u, v):
        """
        接続するノード（接続元u、接続先v）を1つのクラスター番号に統一する機能
        """
        # ノード番号uのクラスター番号を返す
        root_u = self.find(u)
        # ノード番号vのクラスター番号を返す
        root_v = self.find(v)

        # それぞれのクラスター番号が異なる場合、ランク値が高いクラスター番号に揃える
        if root_u != root_v:
            if self.rank[root_u] > self.rank[root_v]:
                self.parent[root_v] = root_u
            elif self.rank[root_u] < self.rank[root_v]:
                self.parent[root_u] = root_v
            else:
                # ランクが同じであれば、どちらかを採用してランク値を増やす
                self.parent[root_v] = root_u
                self.rank[root_u] += 1

    def find_cluster(self,start_node):
        """
        start nodeと同じクラスターに属する全てのノードを返す機能
        """

        cluster = self.find(start_node)
        # 同じクラスターに属するノードをリストアップ
        cluster_nodes = [node for node in range(len(self.parent)) if self.find(node) == cluster]
        return cluster_nodes

    def find_all_cluster(self, check_nodes=None):
        """
        self.parentのcluster情報を使って、ノードをgrouping
        """
        from collections import defaultdict
        from tqdm.auto import tqdm
        
        clusters = defaultdict(list)
        
        if check_nodes is None:
            check_nodes = list(range(0,self.n_node))
            
        # 各ノードのクラスター番号でノードをまとめる    
        for node in tqdm(check_nodes):
            root = self.parent[node]
            clusters[root].append(node)
        
        return list(clusters.values())

n_node = np.unique(list(edge_index_filtered)).max() + 1

uf = UnionFind(n_node=n_node)

for u, v in edge_index_filtered:
    uf.union(u, v)

cluster0 = uf.find_cluster(start_node=0)

clusters = uf.find_all_cluster()