Python实现含公共元素的列表聚合归并方法

需求是将所有存在公共元素的子列表归为同一组，示例如下：

输入：

inputs = [['a','b'], ['a','c'], ['b','d'], ['e','f'], ['g','h'], ['i','k'], ['k','l']]

预期输出：

aggregated_output = [['a','b','c','d'],['e','f'],['g','h'],['i','k','l']]

规则说明：只要子列表间存在公共元素就归为同一组，最终输出的分组顺序、各组内元素顺序没有强制要求。

这个问题本质是求解元素的连通分量，用并查集（DSU）实现效率最高，逻辑清晰，完整代码如下：

from collections import defaultdict

def aggregate_common_lists(inputs):
    # 收集所有出现过的元素
    all_elements = set()
    for sub in inputs:
        all_elements.update(sub)
    
    # 初始化并查集
    parent = {ele: ele for ele in all_elements}

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

    def union(x, y):
        rx, ry = find(x), find(y)
        if rx != ry:
            parent[ry] = rx
    
    # 同一子列表内的元素全部连通，执行合并
    for sub in inputs:
        if not sub:
            continue
        first = sub[0]
        for ele in sub[1:]:
            union(first, ele)
    
    # 按连通分量分组
    res_map = defaultdict(list)
    for ele in all_elements:
        res_map[find(ele)].append(ele)
    
    return list(res_map.values())

测试示例：

inputs = [['a','b'], ['a','c'], ['b','d'], ['e','f'], ['g','h'], ['i','k'], ['k','l']]
print(aggregate_common_lists(inputs))
# 输出结果（顺序不固定，符合要求）：[['a','c','b','d'], ['e','f'], ['h','g'], ['k','i','l']]

实现说明：

• 自动兼容空子列表输入，不会抛出运行错误
• 时间复杂度接近线性，和所有子列表的元素总个数正相关，处理大规模数据性能优异
• 输出结果的分组顺序、组内元素顺序不做固定，完全匹配需求规则

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