Pandas groupby, sum rows, and divide sum by number of rows in group

Viewed 826

I have a dataframe:

>>> import pandas as pd
>>>
>>> df = pd.DataFrame({
...     'P': ['P1', 'P1', 'P2', 'P2', 'P2'],
...     'A1': [0,1,2,1,2],
...     'A2': [5,4,1,3,2],
...     'A3': [5,1,3,8,4],
...     'A4': [2,1,3,4,4],
... })
>>> df
    P  A1  A2  A3  A4
0  P1   0   5   5   2
1  P1   1   4   1   1
2  P2   2   1   3   3
3  P2   1   3   8   4
4  P2   2   2   4   4
>>>

For each P I have to sum columns A1-A4. Than devide this sum by number of P rows. For example, number of rows in each P is:

>>> df.groupby('P').size()
P
P1    2
P2    3
dtype: int64
>>>

Sum of all cols is:

>>> df.groupby('P').sum()
    A1  A2  A3  A4
P
P1   1   9   6   3
P2   5   6  15  11
>>>

but as I need sum by rows, I will use:

>>> df.groupby('P').sum().sum(axis=1)
P
P1    19
P2    37
dtype: int64
>>>

Now I have to divide 19/2 (size) and 37/3 in order to get the results that I need. In order to do that, I would prepare the data like this:

>>> pd.concat([df.groupby('P').sum().sum(axis=1), df.groupby('P').size()], axis=1)
     0  1
P
P1  19  2
P2  37  3
>>>

and than I can use apply in order to get the result:

>>> pd.concat([df.groupby('P').sum().sum(axis=1), df.groupby('P').size()], axis=1).apply(lambda row: row[0]/row[1], axis=1)
P
P1     9.500000
P2    12.333333
dtype: float64
>>>

It works, but I have a feeling that I have overcomplicated calculation for getting sum of rows divided by number of rows for each P.

If someone knows better approach I will be grad to hear it. I would like to get rid of at least concat.

3 Answers

This should work:

df.groupby('P').sum().sum(1) / df.groupby('P').size()

Very round about way:

(df.sum(numeric_only = True, axis = 1)
 .groupby(df.P)
 .pipe(lambda df: df.sum()/df.size())
 )

P
P1     9.500000
P2    12.333333
dtype: float64

You can convert each group to numpy ndarray and sum all values in one step:

df.groupby('P').apply(lambda x: x.to_numpy().sum() / len(x))

Output:

P
P1     9.500000
P2    12.333333
Related