Pandas：按组内复杂条件筛选子数据集的高效实现需求

我明白你遇到的问题——用循环处理分组筛选效率太低了，尤其是数据量大的时候。下面给你一个基于Pandas矢量化操作的高效解决方案，完全避开循环，速度会快很多：

核心思路 #

我们需要先识别出符合条件的分组：

• 对每个custom_id分组，检查是否存在一行sending_num为'from'，且该行的上一行note值为'flag'
• 保留所有符合这个条件的分组的全部行

代码实现 #

import pandas as pd
import numpy as np

# 你的原始DataFrame
df = pd.DataFrame({
    'custom_id': ['aa','aa','aa','aa','aa','aa', 'bk', 'bk', 'bk', 'bk','bk', 'dd', 'dd', 'dd', 'dd', 'dd', 'ff', 'ff', 'ff', 'ff', 'ff', 'ff', 'pu', 'pu', 'pu', 'pu'],
    'sending_num': [11, 252, 198, 266, 5317, 'from', 67, 287, 909, 881, 'from', 22, 55, 'from', 376, 98, 901, 126, 22, 381, 867, 'from', 421, 81, 326, 'from'],
    'receiving_num': [900, 11, 252, 198, 266, 5317, 345, 67, 287, 909, 881, 432, 22, 55, 65, 376, 42, 901, 126, 22, 381, 867, 66, 421, 81, 326],
    'note': [np.nan, 'flag', np.nan, np.nan, 'flag', np.nan, 'flag', np.nan, np.nan, np.nan, np.nan, np.nan, 'flag', np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 'flag', np.nan, np.nan, np.nan, np.nan, np.nan]
})

# 1. 找出所有符合条件的分组
valid_groups = df.groupby('custom_id').apply(
    lambda group: ((group['sending_num'] == 'from') & (group['note'].shift(1) == 'flag')).any()
)

# 2. 筛选原DataFrame中属于有效分组的行
filtered_df = df[df['custom_id'].isin(valid_groups[valid_groups].index)]

# 查看结果
print(filtered_df)

代码解释 #

1. groupby('custom_id').apply(...)：对每个分组执行检查逻辑，group['note'].shift(1)会把每个分组内的note列向下移动一行，这样就能拿到'from'行的上一行的note值
2. ((group['sending_num'] == 'from') & (group['note'].shift(1) == 'flag')).any()：判断当前分组是否存在满足条件的行
3. df[df['custom_id'].isin(...)]：用有效分组的ID筛选原数据，保留这些分组的所有行

结果验证 #

运行后，你会得到aa、dd、ff三个分组的全部行，正好符合你要求的筛选规则，而bk、pu分组会被排除。

效率优势 #

这个方法用Pandas的内置分组和矢量化操作替代了Python循环，底层是C语言优化的计算逻辑，在数据量较大时，速度会比循环实现快几个数量级。

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