时间戳

首页 > 解决方案 > Pandas：按最大值分组并在组上求和的最快方法

python - Pandas：按最大值分组并在组上求和的最快方法

问题描述

这就是我想要实现的目标：

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]

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

标签： pythonpandasgroup-bypandas-groupbyaggregate

解决方案

总的来说，我相信您的方法有效，除了一些改进：

# 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]

另一个改进是你可以做lazy groupby：

groups = df.groupby('A')

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

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

df = df[idx]

推荐阅读