How to compute column sum on the basis of other column value in pandas dataframe?

user_8368524 0
Viewed 47
P T1 T2 T3
0 1 2 3
1 1 2 0
2 3 1 2
3 1 0 2

In the above pandas dataframe df, I want to add columns on the basis of the value of column 'P'.

if df['P'] == 0: 0
if df['P'] == 1: T1 (=1)
if df['P'] == 2: T1+T2 (=3+1=4)
if df['P'] == 3: T1+T2+T3 (=1+0+2=3)

In other words, I want to add from T1 to TN if df['P'] == N. How can I implement this with Python code?

pandas dataframe
1 Answers

EDIT:

For sum values by P column create mask by broadcasting np.arange by length of filtered columns by DataFrame.filter, compare by P values and this mask pass to DataFrame.where, last use sum per rows:

np.random.seed(20)
    
c = [f'{x}{i + 1}' for x in ['T','U','V'] for i in range(3)]
df = pd.DataFrame(np.random.randint(4, size=(10,10)), columns=['P'] + c)

arrP = df['P'].to_numpy()[:, None]

for c in ['T','U','V']:
    df1 = df.filter(regex=rf'^{c}')
    df[f'{c}_SUM'] = df1.where(np.arange(len(df1.columns)) < arrP, 0).sum(axis=1)
print (df)
   P  T1  T2  T3  U1  U2  U3  V1  V2  V3  T_SUM  U_SUM  V_SUM
0  3   2   3   3   0   2   1   0   3   2      8      3      5
1  3   2   0   2   0   1   2   2   3   3      4      3      8
2  0   1   2   2   2   0   1   1   3   1      0      0      0
3  3   2   2   2   1   3   2   1   3   2      6      6      6
4  3   1   1   3   1   2   2   0   2   3      5      5      5
5  2   3   2   3   1   1   1   0   3   0      5      2      3
6  2   3   2   3   3   3   2   1   1   2      5      6      2
7  3   2   0   2   1   1   2   2   2   3      4      4      7
8  2   2   1   0   2   2   0   3   3   0      3      4      6
9  2   2   3   2   2   3   2   2   1   1      5      5      3
user_29010020
Related