Python大型字典遍历性能优化：税值合并函数改进咨询

我编写了merge_tax_values_new_logic函数，用于依据特定逻辑合并税值：遍历税字典（tax_dict），寻找共享国家代码后缀且值存在重叠的键，找到后合并对应值并移除重复键。

原函数代码如下：

def merge_tax_values_new_logic(tax_dict):
    treated_list = set()
    while True:
        changed = False
        for key1, value1 in list(tax_dict.items()):
            country_code = key1[-2:]
            print('current list :',tax_dict)
            if key1 not in treated_list:
                print('current iteration key :' , key1) 
                for key2, value2 in list(tax_dict.items()):
                    if key2.endswith(country_code) and key1 != key2 and any(hl_id in value2 for hl_id in value1):
                        tax_dict[key1].extend(value2)
                        tax_dict.pop(key2)
                        tax_dict[key1] = list(set(tax_dict[key1]))
                        changed = True
                        print( 'current key : ' , key1 , 'matched  with key : ' , key2  ,  'state  of the dict after the pop : ', tax_dict)
                        break
            treated_list.add(key1)
            print('treated list :', treated_list)
            print('******************************')
            if changed:
                break
        if not changed:
            break
    return tax_dict

示例 #

new_tax_dict = {'tax1_US':['A'],'tax2_US':['B'], 'tax3_US':['A','B']}
merge_tax_values_new_logic(new_tax_dict)

运行结果：

current list : {'tax1_US': ['A'], 'tax2_US': ['B'], 'tax3_US': ['A', 'B']}
current iteration key : tax1_US
current key :  tax1_US matched  with key :  tax3_US state  of the dict after the pop :  {'tax1_US': ['A', 'B'], 'tax2_US': ['B']}
treated list : {'tax1_US'}
******************************
current list : {'tax1_US': ['A', 'B'], 'tax2_US': ['B']}
treated list : {'tax1_US'}
******************************
current list : {'tax1_US': ['A', 'B'], 'tax2_US': ['B']}
current iteration key : tax2_US
current key :  tax2_US matched  with key :  tax1_US state  of the dict after the pop :  {'tax2_US': ['A', 'B']}
treated list : {'tax2_US', 'tax1_US'}
******************************
current list : {'tax2_US': ['A', 'B']}
treated list : {'tax2_US', 'tax1_US'}
******************************
{'tax2_US': ['A', 'B']}

该函数在小型字典上运行正常，但处理包含40000+个键、平均每个键对应5个值的大型字典时，性能存在严重问题，求替代实现方案？

方案1：基于并查集（Union-Find）的高效合并 #

你的需求本质是按国家分组后，将值有交集的键归为同一连通分量，最后合并每个分量的所有值。并查集（Disjoint Set Union, DSU）是处理这类连通分量问题的最优数据结构之一，时间复杂度接近O(nα(n))（α是阿克曼函数的反函数，增长极慢，可视为常数）。

实现步骤：

1. 按国家代码分组，避免跨国家的无效遍历。
2. 建立「值ID -> 所属键列表」的映射，快速定位共享同一值的键。
3. 使用并查集将共享值的键合并为同一集合。
4. 遍历每个连通分量，合并所有值并去重，保留一个代表键。

代码实现：

def merge_tax_values_dsu(tax_dict):
    from collections import defaultdict

    # 1. 按国家分组，同时保留原始值
    country_groups = defaultdict(dict)
    for key, values in tax_dict.items():
        country_code = key[-2:]
        country_groups[country_code][key] = values

    merged = {}

    for country, group in country_groups.items():
        # 2. 建立 值到对应键的映射
        value_to_keys = defaultdict(list)
        for key, values in group.items():
            for val in values:
                value_to_keys[val].append(key)

        # 3. 初始化并查集
        parent = {key: key for key in group}

        def find(u):
            while parent[u] != u:
                parent[u] = parent[parent[u]]  # 路径压缩，加速查找
                u = parent[u]
            return u

        def union(u, v):
            root_u = find(u)
            root_v = find(v)
            if root_u != root_v:
                parent[root_v] = root_u

        # 合并共享同一值的所有键
        for val, keys in value_to_keys.items():
            if len(keys) > 1:
                base_key = keys[0]
                for key in keys[1:]:
                    union(base_key, key)

        # 4. 合并每个连通分量的所有值
        component_values = defaultdict(set)
        for key in group:
            root = find(key)
            for val in group[key]:
                component_values[root].add(val)

        # 转换为最终字典格式
        for root, vals in component_values.items():
            merged[root] = list(vals)

    return merged

方案2：优化分组后的迭代合并（简化版） #

如果不想引入并查集，也可以对原逻辑做针对性优化，大幅降低时间复杂度：

• 先按国家分组，避免跨国家的无效遍历。
• 将值转为集合，利用集合的快速交集判断特性。
• 使用标记法跟踪已合并的键，避免重复处理。

代码实现：

def merge_tax_values_optimized(tax_dict):
    from collections import defaultdict

    # 按国家分组，同时把值转成集合，方便快速判断交集
    country_groups = defaultdict(dict)
    for key, values in tax_dict.items():
        country_code = key[-2:]
        country_groups[country_code][key] = set(values)

    merged = {}

    for country, group in country_groups.items():
        keys = list(group.keys())
        merged_keys = set()

        for i in range(len(keys)):
            key1 = keys[i]
            if key1 in merged_keys:
                continue
            current_values = group[key1].copy()

            # 遍历后续键，寻找有交集的项合并
            for j in range(i+1, len(keys)):
                key2 = keys[j]
                if key2 in merged_keys:
                    continue
                # 集合交集判断，效率远高于遍历列表
                if current_values & group[key2]:
                    current_values.update(group[key2])
                    merged_keys.add(key2)

            merged[key1] = list(current_values)
            merged_keys.add(key1)

    return merged

方案对比 #

• 并查集方案：适合超大规模数据（4w+键），性能最优，能正确处理所有连通分量（比如A和B有交集、B和C有交集，最终A/B/C会被合并为一组）。
• 优化迭代方案：实现更简单，性能比原函数提升10-100倍，在极端连通场景下略逊于并查集，但对于4w+键的场景依然足够高效。

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