python - 获取每个组的前 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 等于每个组的“cutoff”中指定的截止值。

感谢任何帮助。提前致谢!

Answer 1

使用 GroupBy.apply 的自定义函数和 DataFrame.head并按 iat 选择每组的第一个值 cutoff :

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