Row sums of dataframe with variable column indexes (Python)

user_6495910 0
Viewed 177

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

pandas row sum python
2 Answers

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
user_43333590

Another option, using numpy:


cols = np.arange(len(df.columns))
# build a 2D array
mask = np.tile(cols, (len(df), 1))
# generate booleans by comparing to `df_ind`
mask = mask >= df_ind.to_numpy()
# replace True with `df`
mask = np.where(mask, df, mask)
# convert nulls to zero, and sum along the columns
mask = np.nan_to_num(mask).sum(1)
df.assign(total = mask)

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