优化Pandas中apply与lambda函数的高效实现方案

问题背景 #

需要实现逻辑：按年份返回指定MSA区域内招生人数最多的学区对应的wage值。原方案使用apply+lambda逐行遍历实现，但面对321681×272的大型数据集时计算速度极慢，需用向量化操作替代遍历提升性能。

原示例代码 #

df = pd.DataFrame({'year': [2000, 2000, 2001, 2001],
                    'msa': ['NYC-Newark', 'NYC-Newark', 'NYC-Newark', 'NYC-Newark'],
                  'leaname':['NYC School District', 'Newark School District', 'NYC School District', 'Newark School District'], 
                  'enroll': [100000,50000,110000,60000],
                   'wage': [5,2,7,3] })


def function1(x,y, var):
    '''
    Returns the selected variable's value for school district with largest enrollment in a given year
    '''

    t = df[(df['msa'] == x) & (df['year'] == y)]
    e = pd.DataFrame(t.groupby(['msa',var]).mean()['enroll'])
    return e.loc[e.groupby(level=[0])['enroll'].idxmax()].reset_index()[var]

df['main_city_wage'] = df.apply(lambda x: function1(x['msa'], x['year'], 'wage'), axis = 1)

原示例输出 #

year         msa                 leaname  enroll  wage  main_city_wage
0  2000  NYC-Newark     NYC School District  100000     5               5
1  2000  NYC-Newark  Newark School District   50000     2               5
2  2001  NYC-Newark     NYC School District  110000     7               7
3  2001  NYC-Newark  Newark School District   60000     3               7

优化方案 #

核心是利用Pandas底层优化的向量化分组操作，替代逐行遍历的apply逻辑，大幅提升计算效率。以下提供两种高效实现方式：

方式1：排序+去重（直观稳定） #

通过对msa、year、enroll排序，保留每组第一条记录即为招生最多的学区，再合并回原数据集：

import pandas as pd

# 生成msa+year对应的最大招生学区薪资映射表
max_enroll_wage = df.sort_values(
    by=['msa', 'year', 'enroll'],
    ascending=[True, True, False]  # 按enroll降序，每组第一条就是招生最多的
).drop_duplicates(subset=['msa', 'year'])  # 保留每个msa+year的第一条记录
[['msa', 'year', 'wage']].rename(columns={'wage': 'main_city_wage'})

# 合并映射表到原数据，确保每一行都能匹配到对应值
df = df.merge(max_enroll_wage, on=['msa', 'year'], how='left')

方式2：分组+索引定位（更简洁） #

利用groupby找到每组enroll最大值的索引，直接提取对应记录的wage：

# 获取每个msa+year组中enroll最大的行索引
max_enroll_idx = df.groupby(['msa', 'year'])['enroll'].idxmax()

# 提取对应行的msa、year和wage，作为映射表
max_enroll_wage = df.loc[max_enroll_idx, ['msa', 'year', 'wage']].rename(columns={'wage': 'main_city_wage'})

# 合并到原数据
df = df.merge(max_enroll_wage, on=['msa', 'year'], how='left')

优化后输出 #

两种方式均会得到与原方案一致的输出：

year         msa                 leaname  enroll  wage  main_city_wage
0  2000  NYC-Newark     NYC School District  100000     5               5
1  2000  NYC-Newark  Newark School District   50000     2               5
2  2001  NYC-Newark     NYC School District  110000     7               7
3  2001  NYC-Newark  Newark School District   60000     3               7

性能说明 #

• 原方案的apply逐行遍历本质是Python层面的循环，大数据集下会产生巨大的性能开销；
• 优化后的两种方式均为Pandas底层C实现的向量化操作，在30万行级别的数据集上，计算速度可提升100-1000倍。

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