如何用Pandas-Numpy高效收集符合条件的列值列表？

原代码使用apply逐行遍历，每次都要对整个DataFrame做布尔索引和去重操作，时间复杂度为O(n²)，在数据集较大时性能极差。可以结合Pandas分组+Numpy向量运算实现高效优化，核心思路是按id_group批量处理，避免逐行重复计算。

优化方案步骤 #

• 按id_group分组，对每个组单独处理：
  1. 提取组内所有level和id值，整理为Numpy数组
  2. 对组内的level去重并排序，同时收集每个唯一level对应的所有唯一id
  3. 对排序后的level序列，累积合并对应的id集合，得到每个level阈值对应的所有符合条件的id列表
• 将处理结果映射回原DataFrame的每一行，恢复原顺序

代码实现 #

import pandas as pd
import numpy as np

# 生成测试数据
test = pd.DataFrame({
    'id_group': [10]*10 + [11]*5,
    'level': list(range(0,10)) + list(range(0,5)),
    'id': [i+20 for i in range(10)] + [i+30 for i in range(5)]
})

def process_group(group):
    # 提取组内的level和id数组
    levs = group['level'].to_numpy()
    ids = group['id'].to_numpy()
    
    # 获取组内唯一的level并排序
    unique_levs = np.unique(levs)
    sorted_levs = np.sort(unique_levs)
    
    # 为每个唯一level收集对应的唯一id
    lev_to_ids = {}
    for lev in sorted_levs:
        lev_to_ids[lev] = np.unique(ids[levs == lev]).tolist()
    
    # 累积合并id列表：每个level对应的是所有小于它的level的id集合
    cumulative_ids = {}
    current_ids = []
    for lev in sorted_levs:
        cumulative_ids[lev] = current_ids.copy()
        # 合并当前level的id并去重
        current_ids = list(set(current_ids + lev_to_ids[lev]))
    
    # 处理最小level的情况（结果为空列表）
    min_lev = sorted_levs[0]
    cumulative_ids.setdefault(min_lev, [])
    
    # 为组内每一行匹配对应的id列表
    group['list of ids'] = [cumulative_ids.get(t, []) for t in levs]
    return group

# 按id_group分组处理，合并结果
test_optimized = test.groupby('id_group').apply(process_group).reset_index(drop=True)

print(test_optimized)

性能对比 #

• 原方法：每行都要扫描整个DataFrame，大数据集下耗时随数据量呈平方增长
• 优化方法：每个组仅处理一次，时间复杂度为O(n log n)（主要来自排序和去重），在十万级以上数据集能带来显著的速度提升

额外说明 #

• 如果你的level是连续整数且无重复，可以省略unique步骤进一步提速
• 若业务不需要去重（原代码用了unique()，所以保留该逻辑），可直接累积原始id列表，省去set操作

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