получение лучших n записей для каждой группы, где n отличается для каждой группы

Question

представьте себе, что это df:

group  time                 index  cutoff
A      2018-04-01 05:07:00  0.8    2
A      2018-04-01 05:07:00  0.7    2
A      2018-04-01 05:07:00  0.6    2
B      2018-04-01 05:57:00  0.9    4
B      2018-04-01 05:57:00  0.87   4
B      2018-04-01 05:57:00  0.7    4
B      2018-04-01 05:57:00  0.67   4
B      2018-04-01 05:57:00  0.5    4
B      2018-04-01 05:57:00  0.47   4

Я бы хотел изменить этот df на:

group  time                 index  cutoff
A      2018-04-01 05:07:00  0.8    2
A      2018-04-01 05:07:00  0.7    2
B      2018-04-01 05:57:00  0.9    4
B      2018-04-01 05:57:00  0.87   4
B      2018-04-01 05:57:00  0.7    4
B      2018-04-01 05:57:00  0.67   4

Это означает, что для каждой группировки (группы, времени) упорядочивается по индексу и выбираются только верхние n записей, где n равно значению отсечения, указанному в 'отсечении' для каждой группы.

Любая помощь приветствуется. Заранее спасибо!

Answer 1

Используйте пользовательскую функцию с GroupBy.apply и DataFrame.head и выберите первое значение cutoff для группы с помощью iat:

df = df.groupby('group', group_keys=False).apply(lambda x: x.head(x['cutoff'].iat[0]))
print (df)
  group                 time  index  cutoff
0     A  2018-04-01 05:07:00   0.80       2
1     A  2018-04-01 05:07:00   0.70       2
3     B  2018-04-01 05:57:00   0.90       4
4     B  2018-04-01 05:57:00   0.87       4
5     B  2018-04-01 05:57:00   0.70       4
6     B  2018-04-01 05:57:00   0.67       4