Fill in values between given indices of 2d numpy array

Question

Given a numpy array,

a = np.zeros((10,10))

[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

And for a set of indices, eg:

start = [0,1,2,3,4,4,3,2,1,0]
end   = [9,8,7,6,5,5,6,7,8,9]

how do you get the "select" all the values/range between the start and end index and get the following:

result = [[1, 0, 0, 0, 0, 0, 0, 0, 0, 1],
          [1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
          [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
          [1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
          [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
          [1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
          [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
          [1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
          [1, 0, 0, 0, 0, 0, 0, 0, 0, 1]]

My goal is to 'select' the all the values between the each given indices of the columns.

I know that using apply_along_axis can do the trick, but is there a better or more elegant solution?

Any inputs are welcomed!!

Answer 1

You can use broadcasting -

r = np.arange(10)[:,None]
out = ((start  <= r) & (r <= end)).astype(int)

This would create an array of shape (10,len(start) . Thus, if you need to actually fill some already initialized array filled_arr , do -

m,n = out.shape
filled_arr[:m,:n] = out

Sample run -

In [325]: start = [0,1,2,3,4,4,3,2,1,0]
     ...: end   = [9,8,7,6,5,5,6,7,8,9]
     ...: 

In [326]: r = np.arange(10)[:,None]

In [327]: ((start  <= r) & (r <= end)).astype(int)
Out[327]: 
array([[1, 0, 0, 0, 0, 0, 0, 0, 0, 1],
       [1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
       [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
       [1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [1, 1, 1, 1, 0, 0, 1, 1, 1, 1],
       [1, 1, 1, 0, 0, 0, 0, 1, 1, 1],
       [1, 1, 0, 0, 0, 0, 0, 0, 1, 1],
       [1, 0, 0, 0, 0, 0, 0, 0, 0, 1]])

If you meant to use this as a mask with 1s as the True ones, skip the conversion to int . Thus, (start <= r) & (r <= end) would be the mask.