Python如何高效提取嵌套列表中关联共现的元素分组

给定存储元素两两共现关系的嵌套列表：

results = [[1,2],[2,1],[3,0],[0,3],[3,4],[2,5]]

需要提取所有存在直接/间接关联的元素分组，期望输出格式如下（组内元素顺序无要求）：

result = [[1,2,5],[3,0,4]]

要求实现方案避免大量多层循环，优先保证性能。

这个需求本质是无向图的连通分量查找问题，最高效的实现方式是使用带路径压缩的并查集（Disjoint Set Union, DSU），整体时间复杂度接近线性，无多层嵌套循环，在大数据量下性能远高于暴力匹配方案。

实现思路 #

• 将每个共现对中的两个元素视作无向图中相连的两个节点
• 初始化时每个元素自身作为独立集合
• 遍历所有共现对，将对中两个元素所属的集合合并
• 最终将同一集合下的所有元素归集为一组，即为所求结果

完整代码 #

class DSU:
    def __init__(self):
        self.parent = dict()

    def find(self, x):
        # 路径压缩优化，将查找路径上的节点直接挂到根节点，后续查找速度接近O(1)
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        # 未出现过的节点初始化父节点为自身
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        # 合并两个节点所属集合
        root_x, root_y = self.find(x), self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

def extract_cooccur_groups(pairs):
    dsu = DSU()
    for a, b in pairs:
        dsu.union(a, b)
    # 按根节点归集所有同组元素
    group_map = dict()
    for node in dsu.parent:
        root = dsu.find(node)
        group_map.setdefault(root, []).append(node)
    return list(group_map.values())

# 测试调用
if __name__ == "__main__":
    results = [[1,2],[2,1],[3,0],[0,3],[3,4],[2,5]]
    output = extract_cooccur_groups(results)
    print(output)
    # 输出结果为 [[1,2,5],[3,0,4]] （组内元素顺序不固定，符合要求）

方案特点 #

• 性能优异：带路径压缩的并查集单次操作时间复杂度为反阿克曼函数级别，可视为常数时间，哪怕处理百万级别的共现对也能快速完成
• 鲁棒性强：输入中存在对称重复对（比如同时存在[1,2]和[2,1]）时不会干扰结果，无需提前对输入做去重处理
• 无额外依赖：纯Python实现，不需要安装第三方库
• 逻辑简洁：全程最多两层平级遍历，无深层嵌套循环

内容的提问来源于stack exchange，提问作者DevPy