Python：替代嵌套循环，对DataFrame分组批量执行函数的优化方案

现有DataFrame结构 #

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

自定义函数示例 #

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['value'].values, 2)
   df3['moving_av'] = 0
   df3.iloc[1:1+len(moving_av), df3.columns.get_loc('moving_av')] = moving_av
   df4 = pd.concat([df4, df3])

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

期望输出结果 #

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

优化的Pythonic实现方法 #

直接使用Pandas的groupby机制替代嵌套循环，代码更简洁且性能更优：

import pandas as pd
import numpy as np

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):
    # 计算移动平均值
    moving_av = running_mean(group['value'].values, 2)
    # 初始化moving_av列并填充对应位置的值
    group['moving_av'] = 0
    group.iloc[1:1+len(moving_av), group.columns.get_loc('moving_av')] = moving_av
    return group

# 按TypeA和TypeB分组处理
result_df = df.groupby(['TypeA', 'TypeB'], group_keys=False).apply(process_single_group)

说明 #

• groupby(['TypeA', 'TypeB']) 直接按指定的两个列进行分组，无需手动遍历唯一值
• apply(process_single_group) 对每个分组应用自定义处理函数，自动将结果合并为完整的DataFrame
• 相比嵌套循环，该方法利用Pandas内部的向量化优化，在数据量较大时性能提升明显，代码可读性也更强

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