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, e.g.:

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.

• wwii: Given an arbitrary ndarray of shape (10,10) how would you use this to select the desired items/elements?
• Divakar: @wwii Well the output ncols = len(start) and output nrows user has to specify. But for a "squeezed in" version, it would be end.max()+1. I have kept it as np.arange(10)[:,None] with 10 being the no. of rows. But yes, in a generic sense, we would need to initialize and assign with those number of rows and cols as the slice ends along each axis.
• Divakar: @wwii On that term, "select", I think OP meant "select and assign". But again, OP could have meant to use the 1s as a selection mask to select from some other array. That's a possible way to interpret as well.
• snowflake: @Divakar Yup, i intended to use it as a mask
• Divakar: @snowflake Do (a*mask).sum(0) or np.einsum('ij,ij->j',a,mask).