How to assign values based on array of to and from indices without a loop?

Viewed 93

Suppose I have a 5 by 5 matrix:

matrix = np.zeros(shape=(5,5))

And a bunch of indices in to and from format:

top = np.array([3,2,3,4,2])
bot = np.array([0,1,2,1,0])

How to assign a value to all of those index ranges at once?

# here is some pseudo code of what I'm trying to accomplish:
# matrix[bot:top] = 1, or matrix[something] = 1

Since top indices and bot indices are arrays I cannot use bot:top slicing.

I could solve this with a loop but this would give suboptimal performance when I have millions of values:

for i, (t,b) in enumerate(zip(top,bot)):
    matrix[b:t,i] = 1

How to do this without looping?

Edit: Here is the expected result after assigning the value.

array([[1., 0., 0., 0., 1.],
       [1., 1., 0., 1., 1.],
       [1., 0., 1., 1., 0.],
       [0., 0., 0., 1., 0.],
       [0., 0., 0., 0., 0.]])

As you can see first column has rows from 0 to 3 changed to 1 and so on.

1 Answers

You may try something like

rows = []
cols = []
for i, (b, t) in enumerate(zip(bot, top)):
    rows.extend(range(b, t))
    cols.extend([i] * (t - b)) 

rows = np.array(rows)
m = np.zeros((5, 5))
m[rows, cols]  = 1.0

but I think there has to be some loops in the generation of the indices.

Related