Pandas: Fastest way to group by max and summing over the group

In general I believe your approach works, except for a few improvements:

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

Another improvement is that you can do lazy groupby:

groups = df.groupby('A')

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

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

df = df[idx]

 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

Try this:

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

It seems the fastest solution so far is provided by wwnde.

My contribution (faster than the original approach, but slower than others):

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

that could be made faster by using @Quang Hoang tip:

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

Benchmark

# 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()

Output:

                 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

This should be faster:

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