如何用向量化方式获取Pandas DataFrame中连续活动对对应的ID？

嘿，这个需求挺常见的，用Pandas的向量化方法完全可以高效解决，不用写慢腾腾的循环。下面给你两种方案，其中第一种是最推荐的全向量化实现，性能最优：

先构造示例数据 #

首先我们先把问题里的DataFrame构造出来，方便演示：

import pandas as pd

df = pd.DataFrame({
    'ID': [1,1,1,1,2,2,3,3],
    'act': ['A','B','C','D','A','B','A','C']
})

方案一：全向掩码+分组筛选（推荐） #

核心思路是在每个ID组内，标记出所有"当前行是A且下一行是B"的情况，然后筛选出存在该情况的ID：

# 生成布尔掩码：当前行act为A，且同ID的下一行act为B
mask = (df['act'] == 'A') & (df.groupby('ID')['act'].shift(-1) == 'B')

# 提取满足条件的ID，去重后转成列表
result = df[mask]['ID'].unique().tolist()

print(result)  # 输出: [1,2]

解释一下：

• df.groupby('ID')['act'].shift(-1)会在每个ID分组内将act列向上偏移一位，这样每个位置的值就对应同组下一行的act内容
• 把"当前行是A"和"下一行是B"两个条件做与运算，得到所有符合"A紧跟B"的行
• 最后取出这些行的ID，去重后就是我们要的结果

方案二：分组结合向量化检查 #

如果习惯用分组函数，也可以用filter结合向量化逻辑实现，性能略逊于方案一但可读性也不错：

def has_ab_sequence(group):
    # 检查组内是否存在A紧跟B的连续序列
    return ((group['act'].iloc[:-1] == 'A') & (group['act'].iloc[1:] == 'B')).any()

# 筛选出符合条件的分组，提取ID去重转列表
result = df.groupby('ID').filter(has_ab_sequence)['ID'].unique().tolist()

print(result)  # 输出: [1,2]

这里的group['act'].iloc[:-1]和group['act'].iloc[1:]分别取每组除最后一行和除第一行的act列，然后逐位对比是否是A和B的组合，只要存在任意一组满足就保留该ID。

两种方法都是向量化实现，避免了Python层面的循环，在数据量较大时优势非常明显。

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