Pandas：如何选择要在groupby结果中显示的列？

Question

我有一个包含10列的数据框，我想从中列出一些行的列，其中ROUGE_L是按Model分组的最大值，我尝试了：

sdf = df.groupby(['Model','Checkpoint'])['ROUGE_L'].max()
 

它打印：

Model  Checkpoint    ROUGE_L
4      1005100       0.204
       1010200       0.202
       1015300       0.205
       1020400       0.203
       1025500       0.204
                     ...  
16000  1030600       0.396
       1035700       0.396
       1040800       0.408

但我希望：

Model  Checkpoint    ROUGE_L
4      1005300       0.205
16000  1040800       0.408

在类似的问题中，我没有发现这样的陈述

Answer 1

如果我们正在为每个模型组寻找max ROUGE_L的检查点，那么：

df.set_index('ROUGE_L').sort_index().reset_index().groupby(['Model'], as_index=False).last()

输出

   Model    Checkpoint  ROUGE_L
2   4       1015300     0.205
7   16000   1040800     0.408

Answer 2

需要更精确的原始数据帧，但下面的代码应该可以工作：

>>> df.loc[df.groupby("Model")["ROUGE_L"].idxmax()]
   Model  Checkpoint  ROUGE_L
2      4     1015300    0.205
7  16000     1040800    0.408

要选择列，请在上一条指令中追加[["Model", "Checkpoint", "ROUGE_L"]]：

>>> df.loc[df.groupby("Model")["ROUGE_L"].idxmax()][["Model", "Checkpoint", "ROUGE_L"]]
   Model  Checkpoint  ROUGE_L
2      4     1015300    0.205
7  16000     1040800    0.408

Answer 3

如果在同一组中有多个最大值，并且希望检索所有最大值，则可以使用

mask = df.groupby('Model')['ROUGE_L'].apply(lambda group: group == group.max())
df_ = df[mask]

print(df_)

   Model  Checkpoint  ROUGE_L
2      4     1015300    0.205
5      4     1025500    0.205
8  16000     1040800    0.408