如何基于stress最大值索引切片DataFrame并保留原行顺序

我有一个包含strain和stress列的DataFrame，需要按以下规则切片：找到stress列的最大值对应的行，然后选取该行及之后的连续5行（注意不是按stress降序排序取前5行）。

当前错误实现 #

现有代码通过对stress列降序排序后取前5行，这会打乱原数据的行顺序，不符合需求：

import pandas as pd

df = pd.DataFrame({"strain": [1,2,4,6,2,4,7,4,8,3,4,7,3,3,6,4,7,4,3,2],
                    "stress": [0,0.2,0.5,0.8,0.7,1,0.7,0.6,0.7,0.8,0.4,0.2,0,-0.5,-0.8,-1,-0.8,-0.9,-0.7,-0.6]})

# Sort by stress values
new_df = df.copy()

new_df = new_df.sort_values(by = ['stress'], ascending = False)

new_df = new_df[0:5]

当前错误输出 #

strain  stress
5       4     1.0
3       6     0.8
9       3     0.8
4       2     0.7
6       7     0.7

正确实现方案 #

核心思路：先定位stress最大值所在的索引，再基于该索引选取连续的5行（包含当前行）：

import pandas as pd

df = pd.DataFrame({"strain": [1,2,4,6,2,4,7,4,8,3,4,7,3,3,6,4,7,4,3,2],
                    "stress": [0,0.2,0.5,0.8,0.7,1,0.7,0.6,0.7,0.8,0.4,0.2,0,-0.5,-0.8,-1,-0.8,-0.9,-0.7,-0.6]})

# 找到stress列最大值对应的索引
max_stress_idx = df['stress'].idxmax()

# 选取最大值所在行及之后的连续5行
new_df = df.loc[max_stress_idx : max_stress_idx + 4]

print(new_df)

预期输出 #

strain  stress
5       4     1.0
6       7     0.7
7       4     0.6
8       8     0.7
9       3     0.8

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