Python中按TypeA&TypeB组合对DataFrame重复执行函数的优化方法

问题背景 #

现有如下结构的Pandas DataFrame df：

TypeA TypeB timepoint value
A     AB     1         10
A     AB     2         10
A     AC     1         5
A     AC     2         15
A     AC     3         10
...
D     DB     1         1
D     DB     2         1

需要针对TypeA和TypeB的唯一组合重复执行自定义函数（示例为滑动平均函数running_mean），并将结果存入新列，替代低效的嵌套for循环方案，同时要求方案适用于更复杂的自定义函数。

示例自定义函数：

import numpy as np
def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return (cumsum[N:] - cumsum[:-N]) / float(N)

原低效嵌套循环实现（存在列名拼写错误、手动拼接易出错等问题）：

df4 = pd.DataFrame()
for i in df['typeA'].unique().tolist():
  df2 = df[df['typeA'] == i]
  for j in df2['typeB'].unique().tolist():
   df3 = df2[df2['typeB'] == j]
   moving_av = running_mean(df3['Wert'].values, 2)
   df3.iloc[1:1+len(moving_av), df3.columns.get_loc('moving_av')] = moving_av
   df5 = pd.concat([df5, df3])

df = pd.merge(df, df5, how='left', on=['typeA', 'Type', 'Kontonummer', 'timepoint'])

期望输出：

TypeA TypeB timepoint value moving_av
A     AB     1         10     NaN
A     AB     2         10     10
A     AC     1         5      NaN
A     AC     2         15     10
A     AC     3         10     12.5
...
D     DB     1         1      NaN
D     DB     2         1      1

通用高效解决方案 #

使用Pandas原生的groupby+apply组合，这是针对分组批量处理的最优方案，既高效又能适配任意复杂的自定义函数。

步骤说明 #

1. 按TypeA和TypeB对原DataFrame分组
2. 定义一个分组处理函数：接收单个分组的子DataFrame，执行自定义逻辑后返回带新列的子DataFrame
3. 用apply将处理函数应用到所有分组，自动合并结果

示例代码 #

import numpy as np
import pandas as pd

# 示例自定义函数（可替换为任意复杂函数）
def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return (cumsum[N:] - cumsum[:-N]) / float(N)

# 定义分组处理逻辑，这是通用的核心部分
def process_single_group(group, window_size=2):
    # 调用自定义函数计算结果
    calculated_values = running_mean(group['value'].values, window_size)
    # 初始化新列为NaN
    group['moving_av'] = np.nan
    # 将计算结果赋值到对应位置（从第window_size个行开始，索引从0则是window_size-1）
    target_indices = slice(window_size-1, window_size-1 + len(calculated_values))
    group.loc[target_indices, 'moving_av'] = calculated_values
    return group

# 应用分组处理
final_df = df.groupby(['TypeA', 'TypeB'], group_keys=False).apply(process_single_group)

方案优势 #

• 高效性：groupby是Pandas优化后的分组机制，远快于手动嵌套循环遍历唯一值
• 通用性：只需修改process_single_group函数内的逻辑，即可适配任意复杂需求（比如拟合模型、自定义统计计算等）
• 简洁性：避免手动拼接DataFrame和列名匹配的错误，代码更易维护

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