如何在Python中基于另一个DataFrame的ID列过滤DataFrame？

要高效过滤df2，仅保留ID存在于df中的行，最优方案是先提取df中ID的唯一值，再用isin()进行过滤——这一方法针对40万行的大数据集也能保持出色性能，因为unique()会缩减匹配范围，isin()本身是pandas优化过的向量化操作。

代码实现 #

import pandas as pd

df = pd.DataFrame({'ID': ['a', 'a', 'b', 'b', 'b', 'b', 'e', 'g'], 
                   'Price': [7, 8, 9, 2, 1, 2, 3, 9]})
df2 = pd.DataFrame({'ID': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'a', 'a', 'a', 'g', 'b', 'b', 'b', 'b', 'e'],
                    'Count': [1, 0, 2, 1, 0, 1, 3, 3, 3, 5, 1, 3, 2, 0, 1, 2]})

# 提取df中ID的唯一值，减少匹配计算量
valid_ids = df['ID'].unique()
# 过滤df2并重置索引
df3 = df2[df2['ID'].isin(valid_ids)].reset_index(drop=True)

print(df3)

输出结果 #

ID  Count
0   a      1
1   b      0
2   e      0
3   g      3
4   a      3
5   a      3
6   a      5
7   g      1
8   b      3
9   b      2
10  b      0
11  b      1
12  e      2

性能说明 #

• 先调用unique()获取df的ID唯一值集合，避免了isin()每次匹配都遍历df的所有行，大幅减少计算开销。
• isin()是向量化操作，比循环或逐行判断快几个数量级，完全适配40万行的数据集规模。

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