Selecting rows from pandas dataframe limited by count per column value

Question

I have a dataframe defined as follows:

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]})

For each class, I'd like to select top n rows where n is specified by count column. The expected output from the above dataframe would be like this:

How can I achieve this?

Answer 1

Use:

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

Output:

   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

You could use

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

Using 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

Here is a solution which groups by class then then looks at the first value in the smaller dataframe and returns the corresponding rows.

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

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

returns

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