如何在Pandas中实现多列自定义聚合（Split-Apply-Combine风格）

给定如下DataFrame：

A x y
foo 0 0
foo 1 1
foo 2 2
foo 3 3
bar 0 2
bar 2 3
bar 4 4
bar 6 5

需求是按列A分组，对每组计算线性回归y=k*x+b，得到结果：

A k b
foo 1.0 0.0
bar 0.5 2.0

尝试用groupby('A')后调用aggregate方法实现：

grouped = table.groupby('A')
grouped.aggregate(f)

def f():
    pass

但发现aggregate会将数据拆分为Series传入函数f，导致无法同时访问x、y两列，需要用Split-Apply-Combine风格实现这类作用于多列的“聚合”功能。

解决方案 #

使用groupby.apply()实现 #

apply方法会将每组的完整DataFrame传递给自定义函数，刚好满足同时访问多列的需求。以下是完整实现代码：

import pandas as pd
import numpy as np

# 构造原始数据
data = {
    'A': ['foo', 'foo', 'foo', 'foo', 'bar', 'bar', 'bar', 'bar'],
    'x': [0, 1, 2, 3, 0, 2, 4, 6],
    'y': [0, 1, 2, 3, 2, 3, 4, 5]
}
table = pd.DataFrame(data)

def calc_lr_params(group):
    # 从组数据中提取x和y的数值数组
    x = group['x'].values
    y = group['y'].values
    # 用numpy.polyfit计算一阶多项式（线性回归）的斜率k和截距b
    k, b = np.polyfit(x, y, 1)
    # 返回Series，方便后续自动合并结果
    return pd.Series({'k': k, 'b': b})

# 分组后应用函数，reset_index将分组索引转为列
result = table.groupby('A').apply(calc_lr_params).reset_index()
print(result)

运行后输出：

A    k    b
0  bar  0.5  2.0
1  foo  1.0  0.0

手动实现Split-Apply-Combine流程 #

如果想更直观地拆分步骤，也可以手动完成三个阶段：

1. Split：按A分组，获取各组的DataFrame
2. Apply：对每组计算线性回归参数
3. Combine：将各组结果合并为最终DataFrame

代码示例：

# Split：拆分数据为各组
groups = dict(table.groupby('A'))

# Apply：遍历每组计算参数
result_list = []
for group_name, group_data in groups.items():
    k, b = np.polyfit(group_data['x'], group_data['y'], 1)
    result_list.append({'A': group_name, 'k': k, 'b': b})

# Combine：合并结果为DataFrame
result = pd.DataFrame(result_list)
print(result)

输出与上一种方法完全一致。

为什么aggregate不适用？ #

aggregate（或其别名agg）的设计逻辑是对每列独立执行聚合操作，比如求每列的均值、总和，所以它会把每组的每一列单独作为Series传入函数，无法同时获取多列数据进行协同计算。而apply是针对整个组的DataFrame进行操作，更适合这类需要多列配合的计算场景。

内容的提问来源于stack exchange，提问作者core_exe