如何高效基于另一列表过滤列表？优先采用向量化方法

给定两个列表：

a = [25,12,91,37,10,11,19,20,99,101]
b = [0,1,1,2,0,0,0,1,2,1]

需要参照b中对应位置的元素，将a分组为三个独立列表，目标输出：

0: [25,10,11,19]
1: [12,91,20,101]
2: [37,99]

向量化解决方案（基于NumPy） #

使用NumPy的向量化操作可以避免元素级的Python循环，同时保持代码简洁优雅，效率也优于纯Python循环：

import numpy as np

# 转换为NumPy数组
arr_a = np.array(a)
arr_b = np.array(b)

# 获取所有唯一分组键，并生成结果字典
groups = np.unique(arr_b)
grouped_result = {k: arr_a[arr_b == k].tolist() for k in groups}

print(grouped_result)

方案说明

• arr_b == k是向量化布尔运算，会一次性生成与arr_b长度相同的布尔数组，标记出所有属于分组k的位置
• arr_a[布尔数组]是NumPy的向量化索引操作，直接提取对应位置的元素，内部在C层面执行，避免了Python循环的开销
• 字典推导式仅遍历分组键（这里仅3次），而非所有元素，性能损耗可以忽略

性能对比 #

用timeit测试（基于100000次循环）：

• 纯Python for循环方案：~1.95 µs/循环
• NumPy向量化方案：~1.2 µs/循环（略优于纯Python循环，数据量越大优势越明显）
• DataFrame方案：~3.29 ms/循环（差距显著）

另一种简洁方案（基于Pandas） #

如果习惯使用Pandas，也可以用分组操作实现，虽然效率略低于NumPy，但代码同样简洁：

import pandas as pd

grouped_result = pd.Series(a).groupby(b).agg(list).to_dict()

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