如何利用Python内置/NumPy函数高效构建元素-索引列表字典？

需求描述 #

给定一个包含N个可哈希不可变对象（如整数对）的列表，其中唯一元素数量约为√N，需要构建一个字典：键为列表中的唯一元素，值为该元素在原列表中所有索引组成的列表。

示例：

l = [(1,2),(1,2),(1,3),(1,4),(5,2)]
d = {(1, 2): [0, 1], (1, 3): [2], (1, 4): [3], (5, 2): [4]}

现有实现（效率待优化） #

当前通过遍历实现，但遍历全部元素的效率有待提升：

def indices_dict(l):
    d = {}
    for i, e in enumerate(l):
        if not e in d:
            d[e] = []
        d[e].append(i)
    return dict(d)

优化方案 #

方案1：使用collections.defaultdict简化并提升效率 #

Python内置的collections.defaultdict可以省去手动判断键是否存在的步骤，内部实现经过优化，能减少循环中的条件分支开销：

from collections import defaultdict

def indices_dict_optimized(l):
    d = defaultdict(list)
    for idx, elem in enumerate(l):
        d[elem].append(idx)
    return dict(d)

这个版本直接为不存在的键自动创建空列表，逻辑更简洁，在元素数量较大时，性能会明显优于原实现。

方案2：使用NumPy实现批量分组（适合大规模数据） #

如果处理的是百万级以上的大规模数据，可借助NumPy的向量化操作提升效率。针对整数对这类元素，我们可以通过结构化数组排序、分组来批量提取索引：

import numpy as np

def indices_dict_numpy(l):
    # 将列表转换为包含元素和原始索引的结构化数组
    arr = np.array(
        [(elem[0], elem[1], idx) for idx, elem in enumerate(l)],
        dtype=[('x', int), ('y', int), ('idx', int)]
    )
    
    # 按元素值排序，让相同元素聚集在一起
    sorted_arr = np.sort(arr, order=['x', 'y'])
    
    # 找出不同元素的分组边界
    _, split_indices = np.unique(sorted_arr[['x', 'y']], return_index=True, axis=0)
    
    # 遍历分组，提取对应索引列表
    result = {}
    for i in range(len(split_indices)):
        start = split_indices[i]
        end = split_indices[i+1] if i+1 < len(split_indices) else len(sorted_arr)
        elem = (sorted_arr[start]['x'], sorted_arr[start]['y'])
        result[elem] = sorted_arr[start:end]['idx'].tolist()
    
    return result

这种方法利用NumPy的底层优化，在超大规模数据场景下效率远超纯Python循环，但小规模数据下，数组转换的额外开销可能会抵消优势，此时优先选择defaultdict版本。

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