如何在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   R

0   C1  A   9

1   C2  A   1

2   C3  A   7

3   C4  B   4

4   C5  B   3

5   C6  B   5

6   C7  B   2

7   C8  C   6

8   C9  C   8

9  C10  C  10

R 列是我的排名列，我想获取排名前 5 的项目（A 列），但是，B 列中每组最多可以选择 3 个项目。

我知道我可以执行以下操作来选择排名前 5 的项目

df.sort_values('R').head(5)

Out[10]: 

    A  B  R

1  C2  A  1

6  C7  B  2

4  C5  B  3

3  C4  B  4

5  C6  B  5

但这会从 B 组中选择 4 个项目。我如何限制它每组最多只能选择 3 个项目？我生成的数据框应该如下所示

    A  B  R

1  C2  A  1

6  C7  B  2

4  C5  B  3

3  C4  B  4

5  C8  C  6

逻辑 - 项目 C6 未被选择，因为它是组 B 的第 4 个项目，因此要选择的下一个可用项目是 C8，它具有下一个最佳排名并且不违反组限制。

Answer 1

我们可以尝试GroupBy.headnew_df = df.sort_values('R').groupby('B', sort=False).head(3).head(5)print(new_df)    A  B  R1  C2  A  16  C7  B  24  C5  B  33  C4  B  47  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   R0   C2  A   11   C3  A   72   C1  A   93   C7  B   24   C5  B   35   C4  B   46   C8  C   67   C9  C   88  C10  C  10