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] 

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

groups = df.groupby('A')

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

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

df = df[idx]

Answer 2

 df.groupby('B').agg(B=('C','max'), C=('D','max'), Sum=('D','sum')).rename_axis('A', axis=0)



    B   C  Sum
A            
s  5  14   27
z  3  12   33

Answer 3

试试这个：

tmp = df.groupby('A').agg(
    idx = ('C', 'idxmax'),
    D = ('D', 'sum')
)
result = df.loc[tmp['idx']].set_index('A').assign(D=tmp['D'])

Answer 4

目前看来最快的解决方案是由 wwnde 提供的。

我的贡献（比原来的方法快，但比其他方法慢）：

df['sum'] = df.groupby('A')['D'].transform('sum')
df = df.loc[df.groupby('A').C.idxmax()]

使用@Quang Hoang 提示可以加快速度：

groups = df.groupby('A')
df['sum'] = groups['D'].transform('sum')
df = df.loc[ groups.C.idxmax()].set_index('A')

基准测试

# Import libraries
import numpy as np
import pandas as pd
from time import time
import seaborn as sns
import matplotlib.pyplot as plt

# Make fake data with 10M rows and 10 target-groups
values = np.arange(10**7)
groups = [f'group{i}' for i in range(1,11) for j in range(int(len(values)/10))]
unused_col = [letter for letter in 'abcdefghij' for j in range(int(len(values)/10))]
df = pd.DataFrame(dict(A=groups, B=unused_col, C=values*0.01, D=values))

# Define functions
def caina_max(df):
    df = df.copy()
    groups = df.groupby('A')
    df['sum'] = groups['D'].transform('sum')
    df = df.loc[ groups.C.idxmax()].set_index('A')
    return df

def Code_Different(df):
    df = df.copy()
    tmp = df.groupby('A').agg(
        idx = ('C', 'idxmax'),
        D = ('D', 'sum'))
    df = df.loc[tmp['idx']].set_index('A').assign(Sum=tmp['D'])
    return df

def Muriel(df):
    df = df.copy()
    df = df.set_index('A')
    df1 = df.groupby(['A','B']).max()
    df2 = df.groupby('A')['D'].sum()
    df = df1.join(df2, lsuffix='_caller', rsuffix='_other')
    df = df.reset_index(level=1).rename(columns={'D_caller': 'D', 'D_other': 'Sum'})
    return df

def Quang_Hoang(df):
    df = df.copy()
    groups = df.groupby('A')
    df['sum'] = groups['D'].transform('sum')
    idx = groups['C'].transform('max') == df['C']
    df = df[idx].set_index('A')
    return df

def valenzio(df):
    df.copy()
    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]
    return df

def wwnde(df):
    df = df.copy()
    df = df.groupby('B').agg(B=('C','max'), C=('D','max'), Sum=('D','sum')).rename_axis('A', axis=0)
    return df

# Benchmark
functions = caina_max, Code_Different, Muriel, Quang_Hoang, valenzio, wwnde
times = {f.__name__: [] for f in functions}

for func in functions:
    fname = func.__name__
    for i in range(100): # reduce this range for faster reproducibility
        t0=time()
        func(df)
        t1=time()
        times[fname].append((t1-t0))

# Benchmark table 
df_benchmark = pd.DataFrame(times).agg([np.mean, np.std, max, min]).T.sort_values('mean').round(3)
df_benchmark.index.name = 'Approach'
# Benchmark figure
plt.figure(figsize=(12,8))
sns.boxplot(data=pd.melt(pd.DataFrame(times)), x='variable', y='value', )
plt.xticks(rotation=45)
plt.title(label='Benchmark', fontweight="bold", pad=20)
plt.ylabel('Time in seconds', labelpad=10)
plt.xlabel('')
plt.show()

输出：

                 mean    std    max    min
Approach                                  
wwnde           1.165  0.009  1.198  1.148
Quang_Hoang     1.488  0.039  1.659  1.439
Code_Different  1.532  0.027  1.638  1.500
caina_max       1.680  0.030  1.813  1.641
valenzio        2.847  0.036  3.030  2.805
Muriel          3.598  0.025  3.666  3.549

Answer 5

这应该更快：

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')

df1 = df.groupby(['A','B']).max()
df2 = df.groupby('A')['D'].sum()
df1.join(df2, lsuffix='_caller', rsuffix='_other')