python - Pandas :按最大值分组和对组求和的最快方法

Question

这是我想要实现的:

input: 
   B  C   D
A          
x  z  1  10
x  z  2  11
x  z  3  12
y  s  4  13
y  s  5  14

output: 
   B  C   D  sum
A               
x  z  3  12   33
y  s  5  14   27

我有以下代码。

import pandas as pd
df = pd.DataFrame({'A': ['x','x','x','y','y'],
               'B': ['z','z','z','s','s'],
               'C': [1,2,3,4,5],
               'D': [10,11,12,13,14]})

df = df.set_index('A') 
df['sum'] = df.groupby('A')['D'].transform('sum')
idx = df.groupby(['A'])['C'].transform(max) == df['C']
df= df[idx]

我正在一个相当大的 Dataframe 上执行此操作。然而这需要很长时间，尤其是第一组。有什么办法可以加快这个过程吗？因为我要做的就是对一组求和并保留不同列最大的行。

Answer 1

总的来说，我相信你的方法是有效的，除了一些改进:

# no need to set_index. Do so on smaller/filtered data if needed
# df = df.set_index('A') 

# this is good 
df['sum'] = df.groupby('A')['D'].transform('sum')

# there's a bit difference between `'max'` and `max`.
# one is vectorized, one is not
idx = df.groupby(['A'])['C'].transform('max') == df['C']

df= df[idx] 

另一个改进是你可以做惰性groupby:

groups = df.groupby('A')

df['sum'] = groups['D'].transform('sum')

idx = groups['C'].transform('max') == df['C']

df = df[idx]