如何高效构建反向查找字典？大数据量场景优化方案

我有一个这样的字典：

old_dict = {'a':[0,1,2], 'b':[1,2,3]}

想要得到一个反向映射的新字典：新字典的键是原字典里的值，值是对应的原字典键的列表，即：

new_dict = {0:['a'], 1:['a','b'], 2:['a','b'], 3:['b']}

目前我用下面的代码实现，但实际场景里我有约10000个原键、100000个新键，这段代码的性能很差：

import numpy as np

# 获取新字典的所有键
new_keys = np.unique(np.hstack([old_dict[key] for key in old_dict]))

# 初始化新字典
new_dict = {key: [] for key in new_keys}
# 遍历每个新键
for new_key in new_keys:
    # 遍历每个原键，检查是否包含当前新键
    for old_key in old_dict:
        if new_key in old_dict[old_key]:
            new_dict[new_key].append(old_key)

想知道有没有更高效的实现方式？比如基于树的算法？也可以用其他更合适的数据类型，目前我在查字典反向查找的资料，还尝试用geopandas的sindex实现。

你当前代码的核心性能瓶颈是反向遍历逻辑：先遍历10万级的新键，再逐个检查1万级的原键，时间复杂度是O(M*N)（M是新键数，N是原键数），数据量大时会产生上亿次无效检查。

下面是几种更高效的实现思路：

1. 正向遍历原字典（最优基础方案） #

直接遍历原字典的每个键值对，把每个值和对应的原键关联起来，时间复杂度是O(K)（K是所有值的总个数），比如每个原键对应10个值，总操作数只有10万次，比原来的100亿次操作快好几个数量级。

普通字典实现 #

old_dict = {'a':[0,1,2], 'b':[1,2,3]}
new_dict = {}

for old_key, values in old_dict.items():
    for val in values:
        # 如果值不在新字典中，先初始化空列表
        if val not in new_dict:
            new_dict[val] = []
        new_dict[val].append(old_key)

用collections.defaultdict简化代码 #

可以用Python标准库的defaultdict省去初始化空列表的判断：

from collections import defaultdict

old_dict = {'a':[0,1,2], 'b':[1,2,3]}
new_dict = defaultdict(list)

for old_key, values in old_dict.items():
    for val in values:
        new_dict[val].append(old_key)

# 如果需要转成普通字典（可选）
new_dict = dict(new_dict)

2. 基于Pandas的批量处理（超大数据量场景） #

如果你的值总数量特别大（比如超过千万级），用Pandas的向量化分组操作会比纯Python循环更高效：

import pandas as pd

old_dict = {'a':[0,1,2], 'b':[1,2,3]}
# 把原字典展开成二维表格
df = pd.DataFrame([(old_key, val) for old_key, vals in old_dict.items() for val in vals],
                  columns=['old_key', 'new_key'])
# 按new_key分组，聚合对应的old_key列表
new_dict = df.groupby('new_key')['old_key'].apply(list).to_dict()

关于树结构或空间索引的说明 #

你提到的基于树的算法、geopandas的sindex其实不太适合这个场景——这类工具主要用于空间范围查询、多维数据检索，而你的需求是简单的键值反向映射，用上述正向遍历的方式已经是最优解，不需要引入复杂的数据结构。

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