Sum values of columns that start with the same text string

Question

I want to take the sum of values (row-wise) of columns that start with the same text string. Underneath is my original df with fails on courses.

Original df:

ID  P_English_2  P_English_3  P_German_1   P_Math_1  P_Math_3  P_Physics_2  P_Physics_4
56            1            3           1          2         0            0            3
11            0            0           0          1         4            1            0
6             0            0           0          0         0            1            0
43            1            2           1          0         0            1            1
14            0            1           0          0         1            0            0

Desired df:

ID  P_English   P_German   P_Math   P_Physics
56          4          1        2           3
11          0          0        5           1 
6           0          0        0           1 
43          3          1        0           2
14          1          0        1           0

Tried code:

import pandas as pd  

df = pd.DataFrame({"ID": [56,11,6,43,14], 
             "P_Math_1": [2,1,0,0,0], 
          "P_English_3": [3,0,0,2,1],
           "P_English_2": [1,0,0,1,0], 
             "P_Math_3": [0,4,0,0,1], 
          "P_Physics_2": [0,1,1,1,0],
           "P_Physics_4": [3,0,0,1,0], 
           "P_German_1": [1,0,0,1,0]}) 

print(df)  

categories = ['P_Math', 'P_English', 'P_Physics', 'P_German'] 

def correct_categories(cols):
     return [cat for col in cols for cat in categories if col.startswith(cat)] 

result = df.groupby(correct_categories(df.columns),axis=1).sum()
 print(result)

score 5 · Accepted Answer · answered Jan 19 '21 at 16:46

Let's try groupby with axis=1:

# extract the subjects
subjects = [x[0] for x in df.columns.str.rsplit('_',n=1)]

df.groupby(subjects, axis=1).sum()

Output:

   ID  P_English  P_German  P_Math  P_Physics
0  56          4         1       2          3
1  11          0         0       5          1
2   6          0         0       0          1
3  43          3         1       0          2
4  14          1         0       1          0

Or you can use wide_to_long, assuming ID are unique valued:

(pd.wide_to_long(df, stubnames=categories,
               i=['ID'], j='count', sep='_')
  .groupby('ID').sum()
)

Output:

    P_Math  P_English  P_Physics  P_German
ID                                        
56     2.0        4.0        3.0       1.0
11     5.0        0.0        1.0       0.0
6      0.0        0.0        1.0       0.0
43     0.0        3.0        2.0       1.0
14     1.0        1.0        0.0       0.0

just a side note that usually, the groupby is faster since everything is done in one go, compared to first doing the transformation, then grouping — sammywemmy, Jan 19 '21 at 19:42

Sum values of columns that start with the same text string

1 Answers1