按多列分组并根据列值提取前 x 行

Question

Groupby based on 'month', 'userid' columns and extract only the values as given in 'cntr' column while deleting the rest.

Input Dataframe is:

import pandas as pd
data = {'month': ['Jan','Jan','Jan','Jan','Feb','Feb','Feb','Feb','Feb','Feb'],
   'userid': ['2345','2345','2345','2345', '2345','2345','2345','5678', '5678', '5678'],
   'cntr': ['3','3','3','3','1','1','1','2','2','2']}

df = pd.DataFrame(data = data, columns = ['month','userid','cntr'])
print(df)

Looks like this,

     month    userid     cntr
0     Jan        2345     3
1     Jan        2345     3 
2     Jan        2345     3
3     Jan        2345     3
4     Feb        2345     1
5     Feb        2345     1
6     Feb        2345     1
7     Feb        5678     2  
8     Feb        5678     2
9     Feb        5678     2

Output required

     month    userid     cntr
0     Jan        2345     3
1     Jan        2345     3 
2     Jan        2345     3
3     Feb        2345     1
4     Feb        5678     2  
5     Feb        5678     2

Answer 1

Use custom lambda function in GroupBy.apply with DataFrame.head :

#if necessary convert to integers
df['cntr'] = df['cntr'].astype(int)

f = lambda x: x.head(x['cntr'].iat[0])
df = df.groupby(['month','userid'], sort=False).apply(f).reset_index(drop=True)
print (df)
  month userid  cntr
0   Jan   2345     3
1   Jan   2345     3
2   Jan   2345     3
3   Feb   2345     1
4   Feb   5678     2
5   Feb   5678     2

Answer 2

You can simply do this using Groupby.head :

In [3446]: df = df.groupby(['month','userid']).head(df.cntr.astype(int))

In [3447]: df
Out[3447]: 
  month userid cntr
0   Jan   2345    3
1   Jan   2345    3
2   Jan   2345    3
4   Feb   2345    1
7   Feb   5678    2
8   Feb   5678    2