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
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