如何在Pandas中按组对多列应用自定义非线性函数

问题分析 #

你遇到的核心问题是：使用groupby时设置了as_index=False，导致apply返回的结果结构与原DataFrame索引无法对齐，最终赋值时出现异常填充。此外，原代码中DataFrame创建存在小错误（col2引用未定义的col1）。

修正后的代码 #

首先修正DataFrame的创建，然后采用正确的分组应用方式：

import pandas as pd
import numpy as np

groups = ['Alice']*2 + ['Ben']*4 + ['Charlie']*3
# 修正col2的赋值逻辑，避免引用未定义变量
df = pd.DataFrame({
    'name': groups,
    'col1': np.arange(len(groups)),
    'col2': 3 * np.arange(len(groups))
})

def f(x,y):
    return np.max(x) + np.max(y) + np.array(x)

# 正确实现：去掉as_index=False，用explode展开分组结果
df['col3'] = df.groupby('name')[['col1','col2']].apply(lambda g: f(g['col1'], g['col2'])).explode()

关键说明 #

1. 去掉as_index=False：保留默认的分组索引方式，让apply返回的每个组的结果以数组形式作为Series的元素，而非DataFrame的行。
2. 使用explode()：将每个组对应的数组元素展开为单独的行，与原DataFrame的每一行索引一一对应，确保赋值后col3的结果符合预期。

执行后col3的结果为：[4, 5, 22, 23, 24, 25, 38, 39, 40]，与预期一致。

另一种实现方式（使用transform逻辑） #

如果函数逻辑允许，也可以用更简洁的方式实现，自动处理索引对齐：

def group_f(g):
    x = g['col1']
    y = g['col2']
    return np.max(x) + np.max(y) + x

df['col3'] = df.groupby('name').apply(group_f).reset_index(drop=True)

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