Python Pandas中带过滤条件的窗口函数等效实现优化

问题背景 #

用户构建了如下测试数据集：

import pandas as pd

data = [['tom', 30, 'sales', 5], ['nick', 35, 'sales', 8], ['juli', 24, 'marketing', 4], ['franz', 40, 'marketing', 6], ['jon', 50, 'marketing', 6], ['jeremie', 60, 'marketing', 6]]
  
df = pd.DataFrame(data, columns=['Name', 'Age', 'Department', 'Tenure'])

需求为：为每行数据计算同部门中年龄大于等于当前行人员年龄的所有人员的平均年龄（例如销售部门的Tom对应平均年龄32.5，Nick对应35）。

用户已有可实现功能的代码，但效率较低，希望找到更优解，现有代码如下：

#Dynamically get mean, where age is greater than the line in question - almost definitely a better
#way of doing this though

def sumWindow(group):

    x = group['Age'].mean()
    group['Mean Dept Age'] = x
    return group

Name = []
Age = []
Department = []
Tenure = []
MeanDeptAge = []

for index, row in df.iterrows():
    
    n = row['Name']
    a = row['Age']
    df_temp = df[df['Age'] >= a]
    df_present = df_temp.groupby(df['Department']).apply(sumWindow)
    df_present['Relevant Name'] = n
    df_final = df_present[df_present['Name'] == df_present['Relevant Name']]
    Name.append(df_final.iloc[0,0])
    Age.append(df_final.iloc[0,1])
    Department.append(df_final.iloc[0,2])
    Tenure.append(df_final.iloc[0,3])
    MeanDeptAge.append(df_final.iloc[0,4])

del df_final    

df_final = pd.DataFrame({'Name': Name, 
                         'Age': Age, 
                         'Department': Department, 
                         'Tenure': Tenure, 
                         'Mean Department Age - Greater Than Emp Age': MeanDeptAge, 
                        })

df_final

用户已尝试多种groupby子句内过滤的方案，仍需更高效的实现方式。

优化实现方案 #

可以通过按部门分组后计算后缀均值的方式实现，彻底避免循环和重复分组操作，大幅提升效率：

import pandas as pd

# 构建测试数据
data = [['tom', 30, 'sales', 5], ['nick', 35, 'sales', 8], ['juli', 24, 'marketing', 4], ['franz', 40, 'marketing', 6], ['jon', 50, 'marketing', 6], ['jeremie', 60, 'marketing', 6]]
df = pd.DataFrame(data, columns=['Name', 'Age', 'Department', 'Tenure'])

# 定义计算后缀均值的函数
def calculate_suffix_mean(group):
    # 按年龄升序排序，确保当前行及之后的行年龄≥当前行
    sorted_group = group.sort_values('Age', ascending=True)
    # 反转年龄列后计算累积均值，再反转回来得到后缀均值（即≥当前年龄的平均年龄）
    sorted_group['Mean Department Age - Greater Than Emp Age'] = sorted_group['Age'][::-1].expanding().mean()[::-1]
    return sorted_group

# 按部门分组应用函数，最后恢复原数据的行顺序
df_final = df.groupby('Department', group_keys=False).apply(calculate_suffix_mean).sort_index()

print(df_final)

方案说明 #

• 分组排序：每个部门内按年龄升序排序，这样当前行及之后的所有行都是年龄≥当前行的员工。
• 后缀均值计算：利用expanding().mean()计算累积均值，通过反转数组的技巧，快速得到每个位置对应的后缀均值（即从当前位置到末尾的平均值）。
• 恢复原顺序：最终按原索引排序，保证结果和原始数据的行顺序完全一致，无需额外调整。

这种方法的时间复杂度远低于原方案，尤其在数据量较大时，性能提升会非常明显。

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