如何在pandas中按组获取最大限制的前n行

Question

我有一个看起来像这样的数据框pd.DataFrame({'A': ['C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9', 'C10'],  ...:                    'B': ['A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'C', 'C'],  ...:                    'R': [9, 1, 7, 4, 3, 5, 2, 6, 8, 10]})Out[3]:      A  B   R0   C1  A   91   C2  A   12   C3  A   73   C4  B   44   C5  B   35   C6  B   56   C7  B   27   C8  C   68   C9  C   89  C10  C  10R 列是我的排名列，我想获取排名前 5 的项目（A 列），但是，B 列中每组最多可以选择 3 个项目。我知道我可以执行以下操作来选择排名前 5 的项目df.sort_values('R').head(5)Out[10]:     A  B  R1  C2  A  16  C7  B  24  C5  B  33  C4  B  45  C6  B  5但这会从 B 组中选择 4 个项目。我如何限制它每组最多只能选择 3 个项目？我生成的数据框应该如下所示    A  B  R1  C2  A  16  C7  B  24  C5  B  33  C4  B  45  C8  C  6逻辑 - 项目 C6 未被选择，因为它是组 B 的第 4 个项目，因此要选择的下一个可用项目是 C8，它具有下一个最佳排名并且不违反组限制。

Answer 1

我们可以尝试GroupBy.head

new_df = df.sort_values('R').groupby('B', sort=False).head(3).head(5)

print(new_df)

    A  B  R

1  C2  A  1

6  C7  B  2

4  C5  B  3

3  C4  B  4

7  C8  C  6

Answer 2

top = df.merge(

    df.groupby('B').R.nsmallest(3) # get the 3 top ranked rows for each group

        .reset_index('B'),

        # `nsmallest` will return a new df with B and df.index as MultiIndex

        # so we reset B to a column

        # however column A is not in this new df, so we merge with the original df

    how='right') # and drop any rows not in the new df

输出

     A  B   R

0   C2  A   1

1   C3  A   7

2   C1  A   9

3   C7  B   2

4   C5  B   3

5   C4  B   4

6   C8  C   6

7   C9  C   8

8  C10  C  10