如何用Python Pandas按列分组后对指定列应用自定义函数？

我有如下的Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': ['one', 'two', 'one', 'three', 'one'],
    'B': [50, 30, 50, 40, 40],
    'C': [35, 40, 35, 35, 35],
    'D': [1.5, 2.0, 3.0, 3.5, 2.5]
})

对应的表格：

我需要先按B、C列分组，再对每组的D列应用一个自定义数学函数，并用计算结果填充新列E。这个自定义函数接收numpy数组作为输入，输出等长数组。比如分组(50,35)对应的D值是[1.5, 3.0]，应用函数后得到对应结果；分组(40,35)对应D值[3.5,2.5]，分组(30,40)对应D值[2.0]。

期望输出：

嗨，这个需求用Pandas的groupby加transform就能完美解决，我给你一步步说清楚：

1. 先定义你的自定义函数 #

首先你得把你的数学逻辑写成一个函数，要求是接收numpy数组，输出长度相同的数组就行。这里我根据你给的示例结果写了个占位函数，你直接替换成自己的实际计算逻辑就好：

def custom_calculation(arr):
    # 这里替换成你的真实计算逻辑
    if len(arr) == 2:
        if arr[0] == 1.5 and arr[1] == 3.0:
            return np.array([4.5, 3.5])
        elif arr[0] == 3.5 and arr[1] == 2.5:
            return np.array([6.8, 8.9])
    elif len(arr) == 1:
        return np.array([4.5])
    # 其他分组的逻辑可以继续补充
    return arr

2. 分组应用函数生成新列E #

核心就是用transform方法，它会自动把分组计算后的结果映射回原DataFrame的每一行，完美保留原来的索引和顺序：

df['E'] = df.groupby(['B', 'C'])['D'].transform(custom_calculation)

为啥用transform而不是apply？ #

简单说：

• apply默认会返回每个分组的聚合结果，结构是按分组来的；而transform会把分组里每个元素的计算结果对应放回原DataFrame的位置，正好符合我们要给每一行加E列的需求。
• 只要你的函数输入输出数组长度一致，不管分组是1行还是多行，transform都能正确处理。

运行完上面的代码，你就能得到你想要的输出结果啦～

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