Adding one column with multiple columns excluding NAN values in Python

I have a dataframe as below:

    FIRST  SECOND  SCORE       S1       S2       S3       S4       S5       S6       S7  S8
0       1       1  38147      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
1       2       1    567  38147.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
2       3       1      0    567.0  38147.0      NaN      NaN      NaN      NaN      NaN NaN
3       4       1      0      0.0    567.0  38147.0      NaN      NaN      NaN      NaN NaN
4       5       1      0      0.0      0.0    567.0  38147.0      NaN      NaN      NaN NaN
5       6       1      0      0.0      0.0      0.0    567.0  38147.0      NaN      NaN NaN
6       7       1      0      0.0      0.0      0.0      0.0    567.0  38147.0      NaN NaN
7       8       1    586      0.0      0.0      0.0      0.0      0.0    567.0  38147.0 NaN
8       1       2   3099      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
9       2       2  30460   3099.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
10      3       2   2372  30460.0   3099.0      NaN      NaN      NaN      NaN      NaN NaN
11      4       2      0   2372.0  30460.0   3099.0      NaN      NaN      NaN      NaN NaN
12      5       2      0      0.0   2372.0  30460.0   3099.0      NaN      NaN      NaN NaN
13      6       2      0      0.0      0.0   2372.0  30460.0   3099.0      NaN      NaN NaN
14      7       2      0      0.0      0.0      0.0   2372.0  30460.0   3099.0      NaN NaN
15      8       2    267      0.0      0.0      0.0      0.0   2372.0  30460.0   3099.0 NaN

I want to add the ‘SCORE‘ column with all the columns starting with ‘S and a number‘ (S1, S2, S3…) such that only the S columns having values gets added up where as the ‘NAN’ remain as ‘NAN’.

The number of postfix after S is already known by a variable ‘repeat = 8‘ which might help in getting the number of columns with S

DESIRED OUTPUT

    FIRST  SECOND  SCORE       S1       S2       S3       S4       S5       S6       S7  S8
0       1       1  38147      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
1       2       1    567  38714.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
2       3       1      0    567.0  38147.0      NaN      NaN      NaN      NaN      NaN NaN
3       4       1      0      0.0    567.0  38147.0      NaN      NaN      NaN      NaN NaN
4       5       1      0      0.0      0.0    567.0  38147.0      NaN      NaN      NaN NaN
5       6       1      0      0.0      0.0      0.0    567.0  38147.0      NaN      NaN NaN
6       7       1      0      0.0      0.0      0.0      0.0    567.0  38147.0      NaN NaN
7       8       1    586    586.0    586.0    586.0    586.0    586.0   1153.0  38733.0 NaN
8       1       2   3099      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
9       2       2  30460  33559.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
10      3       2   2372  30460.0   5471.0      NaN      NaN      NaN      NaN      NaN NaN
11      4       2      0   2372.0  32832.0   3099.0      NaN      NaN      NaN      NaN NaN
12      5       2      0      0.0   2372.0  30460.0   3099.0      NaN      NaN      NaN NaN
13      6       2      0      0.0      0.0   2372.0  30460.0   3099.0      NaN      NaN NaN
14      7       2      0      0.0      0.0      0.0   2372.0  30460.0   3099.0      NaN NaN
15      8       2    267    267.0    267.0    267.0    267.0   2639.0  30727.0   3366.0 NaN

Is there a way to efficiently do it without using for loop so that it works on a big data?

Example input data:

import pandas as pd

data_ex = [
    [1,1,38147,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [2,1,567,38147, np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [3,1,0,567,38147, np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [4,1,0,0,567,38147, np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [5,1,0,0,0,567,38147, np.NaN,np.NaN,np.NaN,np.NaN],
    [6,1,0,0,0,0,567,38147, np.NaN,np.NaN,np.NaN],
    [7,1,0,0,0,0,0,567,38147, np.NaN,np.NaN],
    [8,1,586,0,0,0,0,0,567,38147, np.NaN],
    [1,2,3099,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [2,2,30460,3099,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [3,2,2372,30460,3099,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN],
    [4,2,0,2372,30460,3099,np.NaN,np.NaN,np.NaN,np.NaN,np.NaN ],
    [5,2,0,0,2372,30460,3099,np.NaN,np.NaN,np.NaN,np.NaN,],
    [6,2,0,0,0,2372,30460,3099,np.NaN,np.NaN,np.NaN],
    [7,2,0,0,0,0,2372,30460,3099,np.NaN,np.NaN],
    [8,2,267,0,0,0,0,2372,30460,3099,np.NaN],
       
]

#dataframe
test_data = pd.DataFrame(data_ex, columns = ['FIRST', 'SECOND', 'SCORE', 'S1', 'S2', 'S3', 'S4', 'S5', 'S6', 'S7', 'S8'])

Equivalent R code:

 output<- test_data [, map(.SD, .f = function(.){. + SCORE}),.SDcols = paste0('S',seq(1:repeat))]

Answer

We can use DataFrame.add aligned on axis=0 with the S cols and the SCORE column:

# Get List of S Cols
cols = test_data.columns[test_data.columns.str.match(r'Sd+')]
# Add S cols to SCORE and update
test_data[cols] = test_data[cols].add(test_data['SCORE'], axis=0)

*We don’t need to worry about NaN becoming any other number since NaN + any number is NaN

cols:

Index(['S1', 'S2', 'S3', 'S4', 'S5', 'S6', 'S7', 'S8'], dtype='object')

test_data:

    FIRST  SECOND  SCORE       S1       S2       S3       S4       S5       S6       S7  S8
0       1       1  38147      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
1       2       1    567  38714.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
2       3       1      0    567.0  38147.0      NaN      NaN      NaN      NaN      NaN NaN
3       4       1      0      0.0    567.0  38147.0      NaN      NaN      NaN      NaN NaN
4       5       1      0      0.0      0.0    567.0  38147.0      NaN      NaN      NaN NaN
5       6       1      0      0.0      0.0      0.0    567.0  38147.0      NaN      NaN NaN
6       7       1      0      0.0      0.0      0.0      0.0    567.0  38147.0      NaN NaN
7       8       1    586    586.0    586.0    586.0    586.0    586.0   1153.0  38733.0 NaN
8       1       2   3099      NaN      NaN      NaN      NaN      NaN      NaN      NaN NaN
9       2       2  30460  33559.0      NaN      NaN      NaN      NaN      NaN      NaN NaN
10      3       2   2372  32832.0   5471.0      NaN      NaN      NaN      NaN      NaN NaN
11      4       2      0   2372.0  30460.0   3099.0      NaN      NaN      NaN      NaN NaN
12      5       2      0      0.0   2372.0  30460.0   3099.0      NaN      NaN      NaN NaN
13      6       2      0      0.0      0.0   2372.0  30460.0   3099.0      NaN      NaN NaN
14      7       2      0      0.0      0.0      0.0   2372.0  30460.0   3099.0      NaN NaN
15      8       2    267    267.0    267.0    267.0    267.0   2639.0  30727.0   3366.0 NaN