Row sums of dataframe with variable column indexes (Python)

January 5, 2022

I have a dataframe that has a few million rows. I need to calculate the sum of each row from a particular column index up until the last column. The column index for each row is unique. An example of this, with the desired output, will be:

import pandas as pd

df = pd.DataFrame({'col1': [1, 2, 2, 5, None, 4],
                   'col2': [4, 2, 4, 2, None, 1],
                   'col3': [6, 3, 8, 6, None, 4],
                   'col4': [9, 8, 9, 3, None, 5],
                   'col5': [1, 3, 0, 1, None, 7],
                   })

df_ind = pd.DataFrame({'ind': [1, 0, 3, 4, 3, 5]})

for i in df.index.to_list():
    df.loc[i, "total"] = df.loc[i][(df_ind.loc[i, "ind"]).astype(int):].sum()

print(df)

>>
   col1  col2  col3  col4  col5  total
0   1.0   4.0   6.0   9.0   1.0   20.0
1   2.0   2.0   3.0   8.0   3.0   18.0
2   2.0   4.0   8.0   9.0   0.0    9.0
3   5.0   2.0   6.0   3.0   1.0    1.0
4   NaN   NaN   NaN   NaN   NaN    0.0
5   4.0   1.0   4.0   5.0   7.0    0.0

How can I achieve this efficiently with pandas without using a for loop. Thanks

>Solution :

You can create a like-Indexed DataFrame that lists all of the column positions and then by comparing this DataFrame with df_ind you can create a mask for the entire original DataFrame.

Then mask the original DataFrame and sum to get the row sums based on the appropriate index positions that vary by row.

import pandas as pd

mask = pd.DataFrame({col: df.columns.get_loc(col) for col in df.columns}, 
                    index=df.index)
#   col1  col2  col3  col4  col5
#0     0     1     2     3     4
#1     0     1     2     3     4
#2     0     1     2     3     4
#3     0     1     2     3     4
#4     0     1     2     3     4
#5     0     1     2     3     4

mask = mask.ge(df_ind['ind'], axis=0)
#    col1   col2   col3   col4   col5
#0  False   True   True   True   True
#1   True   True   True   True   True
#2  False  False  False   True   True
#3  False  False  False  False   True
#4  False  False  False   True   True
#5  False  False  False  False  False

df['total'] = df[mask].sum(1)

print(df)
   col1  col2  col3  col4  col5  total
0   1.0   4.0   6.0   9.0   1.0   20.0
1   2.0   2.0   3.0   8.0   3.0   18.0
2   2.0   4.0   8.0   9.0   0.0    9.0
3   5.0   2.0   6.0   3.0   1.0    1.0
4   NaN   NaN   NaN   NaN   NaN    0.0
5   4.0   1.0   4.0   5.0   7.0    0.0