Python字典基于值列表重叠元素合并键的低复杂度实现方案问询

原有代码存在两个核心问题：

• 时间复杂度过高：三层循环嵌套，加上item in val2是对列表的线性查找，整体时间复杂度达到O(k² * n)（k为字典键的数量，n为单列表平均长度），数据规模稍大就会出现明显卡顿
• 逻辑存在缺陷：把正在处理的键key也加入removed集合，会导致后续合并后的列表元素没有被遍历检查，部分重叠场景会漏合并

核心思路是把问题转化为连通分量查找：只要两个键的列表有共享元素，两个键就属于同一个连通分量，最终同一个连通分量的所有列表合并即可。我们可以用并查集（Union-Find）结构快速处理连通性问题，优化后的时间复杂度接近线性。

具体实现步骤 #

• 第一步：构建元素到所属键的映射表，遍历所有元素记录每个元素出现在哪些键的列表中，耗时O(M)，M为所有列表的总元素数
• 第二步：初始化并查集，将共享同一元素的键合并到同一连通集合，每一步并查集操作接近常数时间
• 第三步：按连通分量分组，将同组内所有键的列表拼接，选取组内第一个键作为最终字典的键

优化后代码 #

from collections import defaultdict

class UnionFind:
    def __init__(self, elements):
        self.parent = {e: e for e in elements}
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    def union(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx != fy:
            self.parent[fy] = fx

def consolidate(orig_dict):
    # 构建元素到键的映射
    item_to_keys = defaultdict(list)
    for key, lst in orig_dict.items():
        for item in lst:
            item_to_keys[item].append(key)
    # 初始化并查集
    uf = UnionFind(orig_dict.keys())
    # 合并有共享元素的键
    for keys in item_to_keys.values():
        if len(keys) < 2:
            continue
        main_key = keys[0]
        for k in keys[1:]:
            uf.union(main_key, k)
    # 按连通分量分组合并列表
    groups = defaultdict(list)
    for key in orig_dict.keys():
        root = uf.find(key)
        groups[root].extend(orig_dict[key])
    return dict(groups)

效果验证 #

输入测试用例：

test_dict = {'foo': ["1", "2", "7"], 'bar': ["4", "8", "7"], 'baz': ["5", "6"]}
print(consolidate(test_dict))

输出结果和预期一致：

{'foo': ['1', '2', '7', '4', '8', '7'], 'baz': ['5', '6']}

优化后整体时间复杂度为O(M + k * α(k))，其中α是阿克曼函数的反函数，在k小于1e6时α(k)都小于5，可以近似认为是常数，性能相比原有代码有数量级的提升。

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