Pandas分组筛选指定行：选取满足特定条件的DataFrame数据

现有一个按id和time排序的DataFrame，数据如下：

import pandas as pd
df = pd.DataFrame({'id': [1,1,1,1,2,2,2,2], 
                   'time': [1,2,3,4,1,2,5,6],
                   'is': [0,1,0,0,0,1,0,0]})

对应的表格数据：

需要针对每个id，筛选满足以下至少一个条件的行：

• is == 1
• 位于is == 1的行之后，且与该行的time无间隔（即当前行time = 目标行time + 1）

期望得到的结果：

pd.DataFrame({'id': [1,1,2], 
             'time': [2,3,2],
             'is': [1,0,1]})

对应的表格：

可以通过分组处理结合条件筛选实现，以下提供两种可行方式：

方法一：分组自定义函数（逻辑清晰） #

按id分组后，在每组内标记符合条件的行：

import pandas as pd

# 原始数据
df = pd.DataFrame({'id': [1,1,1,1,2,2,2,2], 
                   'time': [1,2,3,4,1,2,5,6],
                   'is': [0,1,0,0,0,1,0,0]})

def filter_group(group):
    # 获取当前组中is=1的行的time值，加1得到需要匹配的time集合
    target_times = group[group['is'] == 1]['time'] + 1
    # 构建筛选掩码：is=1 或者 time在目标集合中
    mask = group['is'] == 1 | group['time'].isin(target_times)
    return group[mask]

# 分组应用筛选并重置索引
result = df.groupby('id').apply(filter_group).reset_index(drop=True)
print(result)

方法二：向量化操作（大数据集更高效） #

先构建每个id对应的目标time集合，再直接筛选：

import pandas as pd

# 原始数据
df = pd.DataFrame({'id': [1,1,1,1,2,2,2,2], 
                   'time': [1,2,3,4,1,2,5,6],
                   'is': [0,1,0,0,0,1,0,0]})

# 生成每个id对应的is=1行的time+1集合
target_map = df[df['is'] == 1].groupby('id')['time'].apply(lambda x: set(x + 1)).to_dict()

# 构建筛选掩码
mask = df['is'] == 1 | df.apply(lambda row: row['time'] in target_map.get(row['id'], set()), axis=1)

# 筛选并重置索引
result = df[mask].reset_index(drop=True)
print(result)

两种方法运行后都会得到期望的结果：

id  time  is
0   1     2   1
1   1     3   0
2   2     2   1

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