Calculate mean of specific groups in pandas DataFrame (using axis=1) - how to get the specific groups?

Viewed 130

Lets assume that I have DataFrame looking similar to this (but with many more columns and rows):

import pandas as pd
df = pd.DataFrame({'A_1': [1.2, 1.4, 2.2], 
                   'A_2': [1.5, 2.3, 0.2], 
                   'A_3': [2.5, 0.7, 2.0], 
                   'B_1': [1.5, 0.9, 0.6], 
                   'B_2': [0.5, 1.3, 1.2], 
                   'B_3': [1.5, 2.5, 0.5], 
                   'C_1': [1.2, 0.3, 1.2], 
                   'C_2': [2.5, 2.3, 1.2],
                   'C_3': [1.5, 0.4, 0.8]})
df

The result look like this:

    A_1 A_2 A_3 B_1 B_2 B_3 C_1 C_2 C_3
0   1.2 1.5 2.5 1.5 0.5 1.5 1.2 2.5 1.5
1   1.4 2.3 0.7 0.9 1.3 2.5 0.3 2.3 0.4
2   2.2 0.2 2.0 0.6 1.2 0.5 1.2 1.2 0.8

Now I would like to calculate mean using axis=1 but always for specific groups, e.g. A-1, A-2, A-3 and then for group B-1, B-2, B-3, etc. (I mean - for each row but indeed the specific group always). Thus, I tried it with different "for loops" and ".format" function but it doesnt work, example:

ID = ["1","2","3"]
CHAIN = ["A","B","C"]

for CHAINS in CHAIN:
    for IDS in ID:
        df['{}-avg' .format(CHAINS)] = df[['{}_{}' .format(CHAINS,IDS)]].mean(axis=1)    

I also found here that many people use "groupby" function but I dont know how to use it when I want to choose my data by described way.

And my desired output should look like this:

    A_1 A_2 A_3 B_1 B_2 B_3 C_1 C_2 C_3 A-avg       B-avg       C-avg
0   1.2 1.5 2.5 1.5 0.5 1.5 1.2 2.5 1.5 1.733333    1.166667    1.733333
1   1.4 2.3 0.7 0.9 1.3 2.5 0.3 2.3 0.4 1.466667    1.566667    1.000000
2   2.2 0.2 2.0 0.6 1.2 0.5 1.2 1.2 0.8 1.466667    0.766667    1.066667

Can, please, someone suggest how to get it? Thank You.

2 Answers

If want use groupby solution with grouping by all columns names before _ use DataFrame.groupby with lambda function and axis=1 for processing by columns:

df1 = df.join(df.groupby(lambda x: x.split('_')[0], axis=1).mean().add_suffix('-avg'))
print (df1)
   A_1  A_2  A_3  B_1  B_2  B_3  C_1  C_2  C_3     A-avg     B-avg     C-avg
0  1.2  1.5  2.5  1.5  0.5  1.5  1.2  2.5  1.5  1.733333  1.166667  1.733333
1  1.4  2.3  0.7  0.9  1.3  2.5  0.3  2.3  0.4  1.466667  1.566667  1.000000
2  2.2  0.2  2.0  0.6  1.2  0.5  1.2  1.2  0.8  1.466667  0.766667  1.066667

If need select only columns by combinations from lists:

ID = ["1","2","3"]
CHAIN = ["A","B","C"]

from  itertools import product
cols = [f'{b}_{a}' for a, b in product(ID, CHAIN)]
print (cols)

df = df[cols]
print (df)
   A_1  B_1  C_1  A_2  B_2  C_2  A_3  B_3  C_3
0  1.2  1.5  1.2  1.5  0.5  2.5  2.5  1.5  1.5
1  1.4  0.9  0.3  2.3  1.3  2.3  0.7  2.5  0.4
2  2.2  0.6  1.2  0.2  1.2  1.2  2.0  0.5  0.8

df1 = df.join(df.groupby(lambda x: x.split('_')[0], axis=1).mean().add_suffix('-avg'))

An alternative would be to create a mapping of the column names, groupby on the mapping, and concatenate :

mapping = {entry: f"{entry[0]}-avg" for entry in df}
mapping

{'A_1': 'A-avg',
 'A_2': 'A-avg',
 'A_3': 'A-avg',
 'B_1': 'B-avg',
 'B_2': 'B-avg',
 'B_3': 'B-avg',
 'C_1': 'C-avg',
 'C_2': 'C-avg',
 'C_3': 'C-avg'}

#groupby and concat
pd.concat([df, df.groupby(mapping, axis=1).mean()], axis=1)
    A_1 A_2 A_3 B_1 B_2 B_3 C_1 C_2 C_3    A-avg       B-avg    C-avg
0   1.2 1.5 2.5 1.5 0.5 1.5 1.2 2.5 1.5 1.733333    1.166667    1.733333
1   1.4 2.3 0.7 0.9 1.3 2.5 0.3 2.3 0.4 1.466667    1.566667    1.000000
2   2.2 0.2 2.0 0.6 1.2 0.5 1.2 1.2 0.8 1.466667    0.766667    1.066667
Related