User-defined function with arbitrary number of grouping variables in pandas

Admin 1
Viewed 70

My data looks like this:

df = pd.DataFrame({
    'cat_1': ['A'] * 3 + ['B'] * 3,
    'cat_2': ['x', 'y', 'z'] * 2,
    'value': [1, 2, 3, 4, 5, 6]
})

I want to create a function that groups (with a variable number of groups) and sums my data. For instance, the following functions achieve this end for one and two groups, respectively.

def grp_and_sum(data, grp_var, sum_var):

    df = data.groupby([grp_var])[sum_var]\
         .sum()

    return(df)

def grp_and_sum_2(data, grp_var1, grp_var2, sum_var):

    df = data.groupby([grp_var1, grp_var2])[sum_var]\
        .sum()

    return(df)

These functions are more-or-less identical save for the variable number of 'grouping' variables. How do I generalize the first function to accept an arbitrary number of grouping variables? Thank you.

pandas-groupby pandas python
2 Answers

You can use varargs for the groupers, but your sum_var argument will have to be passed in as a keyword argument.

def grp_and_sum_n(data, *args, sum_var):
    return data.groupby([*args])[sum_var].sum()


grp_and_sum_n(df, 'cat_2', sum_var='value')
cat_2
x    5
y    7
z    9
Name: value, dtype: int64

grp_and_sum_n(df, 'cat_1', 'cat_2', sum_var='value')
cat_1  cat_2
A      x        1
       y        2
       z        3
B      x        4
       y        5
       z        6
Name: value, dtype: int64
user_49090870

I'd opt for removing the [] around grp_var. This gives you the full flexibility of the by argument for groupby:

by : mapping, function, label, or list of labels

def grp_and_sum(data, by, sum_var):
    df = data.groupby(by)[sum_var].sum()
    return(df)

grp_and_sum(df, 'cat_1', sum_var='value')
#cat_1
#A     6
#B    15
#Name: value, dtype: int64

grp_and_sum(df, ['cat_1'], sum_var='value')
#cat_1
#A     6
#B    15
#Name: value, dtype: int64

grp_and_sum(df, ['cat_1', 'cat_2'], sum_var='value')
#cat_1  cat_2
#A      x        1
#       y        2
#       z        3
#B      x        4
#       y        5
#       z        6
#Name: value, dtype: int64

Can even pass functions for the index

grp_and_sum(df, lambda x: x//2, sum_var='value')
#0     3
#1     7
#2    11
#Name: value, dtype: int64
user_43333590
Related