将列表输入输出的函数应用于Pandas分组DataFrame出错的排查与修复

问题背景 #

自定义函数用于计算列表总和除以（6×元素值）并返回结果列表：

def my_function(inputList):
  intermediateList = []
  outputList = []
  S = 0
  for x in inputList:
    S += x
    y = 6*x
    intermediateList.append(y)
  for elt in intermediateList:
    z = S / elt
    outputList.append(z)
  return outputList

现有Pandas DataFrame：

import pandas as pd

data = {
  'Class_ID': [1,1,1,1,1,1,1,3,3,3,3],
  'Date': ['1/1/2023','1/1/2023','1/1/2023','1/1/2023','1/1/2023','1/1/2023','1/1/2023','17/4/2022','17/4/2022','17/4/2022','17/4/2022'],
  'Student_ID': [3,4,6,8,9,1,10,5,2,3,4],
  'feature': [1,3,-2,8,4,5,6,0.2,0.1,0.55,0.15]
}

df = pd.DataFrame(data)

执行以下代码后，New_feature列出现NaN和列表值，无法得到每行对应的计算结果：

df['New_feature'] = df.groupby('Class_ID')['feature'].apply(my_function)

错误原因 #

• groupby.apply()接收Python列表作为返回值时，会将整个列表作为单个元素赋值给分组对应的首行，其余行因索引无法与返回值匹配，被填充为NaN。这就是部分行显示列表、其余行是NaN的核心原因。
• 原函数返回列表的设计，不符合Pandas需要的「与分组内每行对应的序列（如Series）」要求，无法完成索引对齐。

修复方案 #

方案1：修改自定义函数返回Pandas Series #

让函数接收Series并返回带原索引的Series，确保Pandas能精准对齐每行数据：

def my_function(inputSeries):
    S = inputSeries.sum()
    # 直接计算每个元素的结果，返回带索引的Series
    return S / (6 * inputSeries)

df['New_feature'] = df.groupby('Class_ID')['feature'].apply(my_function)

方案2：使用transform方法（更简洁） #

transform专为分组后返回与原数据同长度结果的场景设计，自动处理索引对齐：

# 自定义函数版本
def my_transform_func(inputSeries):
    S = inputSeries.sum()
    return S / (6 * inputSeries)

df['New_feature'] = df.groupby('Class_ID')['feature'].transform(my_transform_func)

# 或者用lambda简化写法
df['New_feature'] = df.groupby('Class_ID')['feature'].transform(lambda x: x.sum()/(6*x))

额外优化：避免除以0错误 #

如果feature列存在0值，会触发除以0的异常，可添加判断处理：

def my_function_safe(inputSeries):
    S = inputSeries.sum()
    # 将0值替换为inf，避免除以0报错
    return S / (6 * inputSeries.where(inputSeries != 0, float('inf')))

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