从熊猫数据框中选择受每列计数限制的行

Question

我有一个定义如下的数据框：

df = pd.DataFrame({'id':    [11, 12, 13, 14, 21, 22, 31, 32, 33], 
                   'class': ['A', 'A', 'A', 'A', 'B', 'B', 'C', 'C', 'C'],
                   'count': [2, 2, 2 ,2 ,1, 1, 2, 2, 2]})

对于每个类，我想选择n由计数列指定的前n行。上述数据框的预期输出将是这样的：

我该如何实现？

Answer 1

使用：

(df.groupby('class', as_index=False, group_keys=False)
   .apply(lambda x: x.head(x['count'].iloc[0])))

输出：

   id class  count
0  11     A      2
1  12     A      2
4  21     B      1
6  31     C      2
7  32     C      2

Answer 2

您可以使用

In [771]: df.groupby('class').apply(
                     lambda x: x.head(x['count'].iloc[0])
                  ).reset_index(drop=True)
Out[771]:
   id class  count
0  11     A      2
1  12     A      2
2  21     B      1
3  31     C      2
4  32     C      2

Answer 3

使用cumcount

df[(df.groupby('class').cumcount()+1).le(df['count'])]
Out[150]: 
  class  count  id
0     A      2  11
1     A      2  12
4     B      1  21
6     C      2  31
7     C      2  32

Answer 4

这是一个按 class 分组的解决方案，然后查看较小数据框中的第一个值并返回相应的行。

def func(df_):
    count_val = df_['count'].values[0]
    return df_.iloc[0:count_val]

df.groupby('class', group_keys=False).apply(func)

返回

  class  count  id
0     A      2  11
1     A      2  12
4     B      1  21
6     C      2  31
7     C      2  32