用pandas优化apply和lambda函数

Question

我试图优化一个函数，返回每个year给定条件(MSA内最大注册)的变量值(wage)。我认为结合apply和lambda将是有效的，但我的实际数据集很大(形状为321681x272)使得计算非常慢。有更快的方法吗?我认为向量化操作而不是通过df迭代可能是一种解决方案，但我不确定它将遵循的结构作为df.apply和lambda的替代方案

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_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

Answer 1

类似

df['main_wage'] = df.set_index('wage').groupby(['year', 'msa'])['enroll'].transform('idxmax').values