Effective .transform in python

9/1/19

Merging summary statistics back into a table is quite a common thing to do. At first glance, a solution to this problem is to simply compute the statistics in a new table and merge the table back in. However, with pandas we don’t have to do that.

import pandas as pd
import numpy as np

np.random.seed(124)
df = pd.DataFrame({
    'group':np.random.choice(['a', 'b'], 5),
    'x':np.random.randint(100, 200, 5)
}).sort_values(by='group'); df

Here I’ve got some rows that having a grouping column called group, and I’d like to calculate the sum of x within each group and integrate it into my table.

df['g_sum'] = df.groupby('group')['x'].transform(np.sum); df

A key idea is that .transform will apply its function argument to the dataframe and return a result that is the same size as the input frame.

We can also use .transform with a user defined function

def max_minus_one(col):
    """ Get the max - 1 of a list"""
    return np.max(col) - 1

df['udf'] = df.groupby('group')['x'].transform(max_minus_one); df

…and even lambda functions

df['use_lam'] = df.groupby('group')['x'].transform(lambda x: np.max(x) - 1); df