pandas - 按最大值作为索引对数据帧进行切片

Question

我目前有一个应变和应力的数据框，包含相应的值。我想以特定方式对数据帧进行切片 - 我想找到压力中的最大值，然后获取数据帧的接下来 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]

这是我当前的输出:

print(new_df)

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

所以我的代码是按压力的最高值排序的。但是，我想将行顺序保持在列中最高值后面。这将是我的预期输出:

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

Answer 1

您可以使用argmax找到最大值的索引:

imax = df.stress.argmax()
df.iloc[imax:imax+5]

结果:

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