Sum dataframe columns in a loop based on column id

user_11083050 0
Viewed 308

My input dataframe looks something like this.

EXP 1      EXP 2      EXP 3       EXP 4
 1          2           4          3
 5          6           4          3
 5          2           1          2
 3          3           2          3

I want to create three new columns. The first new column is the sum of EXP 1 and EXP 2 The second new column is the sum of EXP 1 and EXP 2 and EXP 3 The third new column is the sum of EXP 1 and EXP 2 and EXP 3 and EXP 4

I've tried making a loop for this but I'm not really sure how to go about it.

Obviously I could just go df['NEW COLUMN 1'] = df['EXP 1'] + df['EXP 2'] etc but this is too slow for my real objective.

Any help is appreciated.

pandas python
1 Answers

Try expanding sum on axis=1:

sums = df.expanding(axis=1, min_periods=2).sum().iloc[:, 1:]
sums.columns = map(lambda x: f'{df.columns[0]}-{x}', sums.columns)

Then merge back to the original DataFrame:

Via concat on axis=1:

new_df = pd.concat((df, sums), axis=1)

or via join:

new_df = df.join(sums)

new_df:

   EXP 1  EXP 2  EXP 3  EXP 4  EXP 1-EXP 2  EXP 1-EXP 3  EXP 1-EXP 4
0      1      2      4      3          3.0          7.0         10.0
1      5      6      4      3         11.0         15.0         18.0
2      5      2      1      2          7.0          8.0         10.0
3      3      3      2      3          6.0          8.0         11.0

Complete Working Example:

import pandas as pd

df = pd.DataFrame({
    'EXP 1': [1, 5, 5, 3],
    'EXP 2': [2, 6, 2, 3],
    'EXP 3': [4, 4, 1, 2],
    'EXP 4': [3, 3, 2, 3]
})

sums = df.expanding(axis=1, min_periods=2).sum().iloc[:, 1:]
sums.columns = map(lambda x: f'{df.columns[0]}-{x}', sums.columns)
new_df = df.join(sums)

print(new_df)
user_154978880
Related