Sum over columns in dynamic range

Question

I am working making a dynamic table by adding columns in same row, however, range of columns is determined based on two columns difference (high-low):

df = pd.DataFrame({
  '10': [1, 10, 20, 30, 40, 50],
  '20': [20, 15, 12, 18, 32, 12],
  '30': [3, 11, 25, 32, 13, 4],
  '40': [32, 11, 9, 82, 2, 1],
  '50': [9, 5, 11, 11, 2, 5],
  'low': [12, 22, 18, 27, 23, 15],
  'high': [45, 41, 33, 54, 35, 45],
})

df

Index     10      20      30     40    50   low    high
0         1       20      3      32    9    12     45
1         10      15      11     11    5    22     41
2         20      12      25     9     11   18     33
3         30      18      32     82    11   27     54
4         40      32      13     2     2    23     35     
5         50      12      4      1     5    15     45

high and low range is then taken to determine which columns are selected, and finally sums by index. So my initial code starts by determining difference between columns, query the (cols) to be used for the operation

def colrange(first, last):
    return (first - last).abs().argsort()[0]

cols = df.columns[:-2]

Then I used iterrows() to start looking in every row between the range:

c = cols.to_series().astype(int)
for idx,row in df.iterrows():
    df.loc[idx,'result']= row[cols[ colrange(c,row.low) : colrange(c, row.high) ]].sum()

So my df['result'] should look like:

Index     10      20      30     40    50   low    high   result
0         1       20      3      32    9    12     45     1+20+3  = 24
1         10      15      11     11    5    22     41     15+11   = 26
2         20      12      25     9     11   18     33     12      = 12
3         30      18      32     82    11   27     54     32+82     = 114
4         40      32      13     2     2    23     35     32        = 32
5         50      12      4      1     5    15     45     50+12+4   = 66

My problem is that this method is too slow, could you advice any other idea how to solve this exercise? I appreciate any thoughts in advance.

user_10403090 · Answer

This is about 5 times faster on your example. It should also scale pretty good as the DataFrame size increase

start = np.abs((c.to_frame().to_numpy().T - df['low'].to_frame().to_numpy())).argsort()[:, 0]
stop = np.abs((c.to_frame().to_numpy().T - df['high'].to_frame().to_numpy())).argsort()[:, 0]
df['result'] = [*map(lambda first, last, row: df.iloc[row, first:last].sum(), start, stop, range(len(df)))]

Sum over columns in dynamic range

1 Answers