如何在含不同长度列表列的Pandas DataFrame中使用NumPy广播过滤行

你的代码错误在于两点：一是对pandas Series做维度扩展后直接比较列表，触发了形状不匹配的对齐错误；二是比较逻辑是判断列表是否相等，而非检查元素是否存在于目标列表中。以下是针对需求的正确实现，兼顾大数据集的效率：

方法一：利用集合查找（推荐，适合大数据集） #

由于列表的in操作是O(n)，而集合是O(1)，先将col_b的列表转为集合，再检查col_a的元素是否存在：

import pandas as pd
import numpy as np

dt = pd.DataFrame({'id' : ['a', 'a', 'a', 'b', 'b'],
                   'col_a': [[1],[2],[5],[1],[2]],
                   'col_b': [[2],[2,4],[2,5,7],[4],[3,2]],
                   })

# 生成过滤掩码：检查col_a的单个元素是否在col_b的集合中
mask = dt.apply(lambda row: row['col_a'][0] in set(row['col_b']), axis=1)
# 过滤行
filtered_dt = dt[mask]
print(filtered_dt)

运行结果：

id col_a       col_b
1  a   [2]     [2, 4]
2  a   [5]  [2, 5, 7]
4  b   [2]     [3, 2]

方法二：NumPy向量化实现 #

如果想借助NumPy的向量化特性，可以先提取col_a的元素，再结合np.vectorize完成检查：

# 提取col_a中的单个元素为NumPy数组
col_a_vals = dt['col_a'].str[0].to_numpy()
# 将col_b转为集合数组
col_b_sets = dt['col_b'].apply(set).to_numpy()

# 定义向量化检查函数
check_exists = np.vectorize(lambda val, s: val in s)
# 生成掩码
mask = check_exists(col_a_vals, col_b_sets)
# 过滤行
filtered_dt = dt[mask]

说明 #

• 若col_a并非单元素列表，只需将判断逻辑改为bool(set(row['col_a']) & set(row['col_b']))，即可检查两个列表是否有交集。
• 对于包含数百个元素的列表，使用集合能大幅提升查找效率，更适合大数据场景。

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