pandas groupby().head(n) 其中 n 是组标签的函数

小能豆

pandas groupby().head(n) 其中 n 是组标签的函数

我有一个数据框，我想按列分组并取每个组的头，但我希望头的深度由组标签的函数定义。如果不是因为组大小可变，我可以轻松做到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

阅读 11

2024-11-19

共1个答案

小能豆

可以使用 Pandas 的 groupby 和 apply 方法来实现需求，同时避免显式迭代，保持高效性。下面是代码实现：

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

# 使用 groupby 和 apply 根据深度获取每组的头
result = df.groupby('label', group_keys=False).apply(lambda group: group.head(depth(group.name)))

# 打印结果
print(result)

代码解释：

按列分组: groupby('label', group_keys=False) 按 label 列分组。
group_keys=False 确保输出数据框不会额外添加分组标签作为索引。
自定义函数应用于每组: 使用 apply 对每个分组执行 group.head(depth(group.name)) 操作。
group.name 是当前分组的名称（例如 'apple'、'car' 等），传递给 depth 函数以确定每组要获取的行数。
避免显式迭代: 整个操作基于 Pandas 的分组功能，避免了对唯一标签的显式迭代，性能更优。

输出结果：

   label  values
0  apple       7
3    car       9
4    car       6
5    dog       5
6    dog       3
7    dog       2

这个方法不仅简洁，而且充分利用了 Pandas 的优化，适合处理大规模数据集。

2024-11-19