我有一个数据框,我想按列分组并获取每个组的头部,但我希望头部的深度由组标签的函数定义。如果不是可变的组大小,我可以轻松做到df.groupby('label').head(n)
。我可以想象一个涉及迭代的解决方案df['label'].unique()
,切片数据帧并构建一个新的数据帧,但我所处的环境对性能非常敏感,因此我希望尽可能避免这种迭代。
这是一个示例数据框:
label values
0 apple 7
1 apple 5
2 apple 4
3 car 9
4 car 6
5 dog 5
6 dog 3
7 dog 2
8 dog 1
以及我的示例设置的代码:
import pandas as pd
df = pd.DataFrame({'label': ['apple', 'apple', 'apple', 'car', 'car', 'dog', 'dog', 'dog', 'dog'],
'values': [7, 5, 4, 9, 6, 5, 3, 2 ,1]})
def depth(label):
if label == 'apple': return 1
elif label == 'car': return 2
elif label == 'dog': return 3
我想要的输出是一个数据框,其中包含该函数定义的每个组的行数:
label values
0 apple 7
3 car 9
4 car 6
5 dog 5
6 dog 3
7 dog 2