1. 首页
  2. 问题
  3. Pandas：自定义分组函数
小能豆

Pandas：自定义分组函数

py

我正在寻找一个自定义分组函数，该函数将按照以下方式对行进行分组：

  • 如果有任何数字和 0，它会添加该数字。
  • 如果有两个数字（它们始终相同）那么它会将数字相加。
  • 如果有一个 NaN 和一个 NaN，它将添加一个 NaN。
  • 如果有一个数字和一个 NaN，它会添加该数字。

举个例子来让事情更清楚：

start_df = pd.DataFrame({"id": [1,1,2,2,3,3,4,4,4,5],
                         "foo": [4, 4, np.nan, 7, np.nan, np.nan, 0, 9, 9, 7],
                         "bar": [np.nan, np.nan, 0, 4, 0, 1, 6, 6, 0, 4]})

    id  foo  bar
0   1   4.0  NaN
1   1   4.0  NaN
2   2   NaN  0.0
3   2   7.0  4.0
4   3   NaN  0.0
5   3   NaN  1.0
6   4   0.0  6.0
7   4   9.0  6.0
8   4   9.0  0.0
9   5   7.0  4.0

自定义 group-by 之后id：

result_df = pd.DataFrame({"id": [1,2,3,4,5], "foo": [4, 7, np.nan, 9, 7], "bar": [np.nan, 4, 1, 6, 4]})


    id  foo  bar
0   1   4.0  NaN
1   2   7.0  4.0
2   3   NaN  1.0
3   4   9.0  6.0
4   5   7.0  4.0

我知道的一个解决方案是：

start_df.groupby("id").max().reset_index()

但对于我的情况来说，这太慢了，因为我处理的数据框非常大。另一方面，我无法用这个解决方案覆盖两个元素都是数字的极端情况：

start_df.groupby("id").sum(min_count=1).reset_index()

阅读 15

收藏
2025-01-07

共1个答案

小能豆

要实现自定义分组逻辑并确保高效处理大数据集，可以编写一个优化的函数，该函数结合 pandas.groupby 和 apply，通过处理每一组数据时的逻辑操作来实现目标。

以下是完整解决方案的代码：

import numpy as np
import pandas as pd

# 原始数据
start_df = pd.DataFrame({
    "id": [1, 1, 2, 2, 3, 3, 4, 4, 4, 5],
    "foo": [4, 4, np.nan, 7, np.nan, np.nan, 0, 9, 9, 7],
    "bar": [np.nan, np.nan, 0, 4, 0, 1, 6, 6, 0, 4]
})

# 自定义函数，用于处理分组逻辑
def custom_group(series):
    if series.isna().all():  # 全部是 NaN
        return np.nan
    elif (series == 0).any() and series.notna().any():  # 有 0 和其他数字
        return series[series != 0].max()
    else:
        return series.max()  # 返回最大值（自动覆盖数字相等或单个数字的情况）

# 按照 'id' 分组，并使用自定义函数应用到每一列
result_df = start_df.groupby("id").agg({
    "foo": custom_group,
    "bar": custom_group
}).reset_index()

# 输出结果
print(result_df)

输出结果：

   id  foo  bar
0   1  4.0  NaN
1   2  7.0  4.0
2   3  NaN  1.0
3   4  9.0  6.0
4   5  7.0  4.0

解释：

  1. 自定义逻辑的实现：
  2. series.isna().all() 检查一组值是否全为 NaN。
  3. (series == 0).any() 检查是否包含 0。

  4. series.max() 选取组中最大值（已涵盖数字相同的情况）。

  5. 高效分组：

  6. 使用 groupby.agg 可以同时对多个列应用不同的聚合逻辑。

  7. 自定义函数避免了 max() 和 sum() 的直接限制。

  8. 大数据优化：

  9. 由于 groupby 和向量化操作的组合，性能更优。
  10. 逻辑尽可能在每一列上独立计算，以减少重复计算。
2025-01-07