如何在Pandas中不使用for循环获取counter为1的frames起止位置

原始数据 #

创建目标DataFrame的代码：

import pandas as pd
frames = [9,12,14,15,16,17,18,22,23,24,25,30]
counter = [0,0,0,0,1,1,1,0,0,1,1,0]
df = pd.DataFrame({'frames':frames, 'counter':counter})

输出的DataFrame内容：

frames  counter
0        9        0
1       12        0
2       14        0
3       15        0
4       16        1
5       17        1
6       18        1
7       22        0
8       23        0
9       24        1
10      25        1
11      30        0

需求 #

不使用for循环，获取counter值为1的所有连续区间对应的frames的起始和结束位置，预期结果：

[(16,18), (24,25)]

解决方案 #

利用Pandas的分组标记技巧实现：

# 标记连续的counter=1区间
df['group'] = (df['counter'] != df['counter'].shift()).cumsum()

# 筛选counter=1的组，计算每组的起止frames
result = df[df['counter'] == 1].groupby('group')['frames'].agg(['first', 'last']).apply(tuple, axis=1).tolist()

print(result)
# 输出: [(16, 18), (24, 25)]

代码说明 #

• df['counter'] != df['counter'].shift()：对比当前行与上一行的counter值，判断是否进入新的连续区间，返回布尔序列
• .cumsum()：将布尔序列转换为累加的分组编号，同一连续区间会被分配相同编号
• 筛选出counter=1的行后，按分组编号聚合，提取每组frames的首尾值并转为元组列表

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