如何从pandas DataFrame中筛选每个steps序列最大值对应的行

核心思路 #

通过识别每组序列的起始标记（steps值为1）生成分组ID，按ID分组后筛选每组steps最大值对应的行即可。

完整实现代码 #

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame([(234, 1), (234, 1), (234, 2), (234, 3), (234, 4), (234, 1), (234, 2),
                   (234, 3), (234, 1), (234, 2)], columns=['Energies', 'steps'])

# 生成分组标识：每次遇到steps=1则开启新分组
group_ids = (df['steps'] == 1).cumsum()

# 按分组筛选每组steps最大值对应的行
result = df.loc[df.groupby(group_ids)['steps'].idxmax()]

结果验证 #

运行后得到的result索引为[0,4,7,9]，完全符合需求，输出内容如下：

Energies  steps
0       234      1
4       234      4
7       234      3
9       234      2

逻辑说明 #

• df['steps'] == 1生成布尔序列，所有新组的起始位置标记为True
• 调用cumsum()对布尔值累加，同一组的行会得到相同的数值作为分组ID
• 按分组ID分组后，idxmax()会返回每组steps最大值对应的行索引
• 最后通过loc筛选出对应行即可，无需生成额外的临时列

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