GroupBy aggregate function that computes two values at once

user_3767239 0
Viewed 73

I have a datafame like the following:

import pandas as pd

df = pd.DataFrame({
    'A': [1, 1, 1, 2, 2, 2],
    'B': [1, 2, 3, 4, 5, 6],
    'C': [4, 5, 6, 7, 8, 9],
})

Now I want to group and aggregate with two values being produced per group. The result should be similar to the following:

expected = df.groupby('A').agg([min, max])

#     B       C    
#   min max min max
# A                
# 1   1   3   4   6
# 2   4   6   7   9

However, in my case, instead of two distinct functions min and max, I have one function that computes these two values at once:

def minmax(x):
    """This function promises to compute the min and max in one go."""
    return min(x), max(x)

Now my question is, how can I use this one function to produce two aggregation values per group?

It's kind of related to this answer but I couldn't figure out how to do it. The best I could come up with is using a doubly-nested apply however this is not very elegant and also it produces the multi-index on the rows rather than on the columns:

result = df.groupby('A').apply(
    lambda g: g.drop(columns='A').apply(
        lambda h: pd.Series(dict(zip(['min', 'max'], minmax(h))))
    )
)

#        B  C
# A          
# 1 min  1  4
#   max  3  6
# 2 min  4  7
#   max  6  9
pandas-groupby pandas aggregate python
2 Answers

If you are stuck with a function that returns a tuple of values. I'd:

  1. Define a new function that wraps the tuple values into a dict such that you predefine the dict.keys() to align with what you want the column names to be.
  2. Use a careful for loop that doesn't waste time and space.

Wrap Function

# Given Function
def minmax(x):
    """This function promises to compute the min and max in one go."""
    return min(x), max(x)

# wrapped function
def minmax_dict(x):
    return dict(zip(['min', 'max'], minmax(x)))

Careful for loop

I'm aiming to pass this dictionary into the pd.DataFrame constructor. That means, I want tuples of the MultiIndex column elements in the keys. I want the values to be dictionaries with keys being the index elements.

dat = {}
for a, d in df.set_index('A').groupby('A'):
    for cn, c in d.iteritems():
        for k, v in minmax_dict(c).items():
            dat.setdefault((cn, k), {})[a] = v

pd.DataFrame(dat).rename_axis('A')

    B       C    
  min max min max
A                
1   1   3   4   6
2   4   6   7   9

Added Detail

Take a look at the crafted dictionary

data

{('B', 'min'): {1: 1, 2: 4},
 ('B', 'max'): {1: 3, 2: 6},
 ('C', 'min'): {1: 4, 2: 7},
 ('C', 'max'): {1: 6, 2: 9}}
user_23366540

One other solution:

pd.concat({k:d.agg(minmax).set_axis(['min','max'])
           for k,d in df.drop('A',axis=1).groupby(df['A'])
          })

Output:

       B  C
1 min  1  4
  max  3  6
2 min  4  7
  max  6  9
user_42384080
Related