How to fill elements of a matrix with indices taken from another matrix in python

user_9890459 0
Viewed 637

I have the following code, which constructs two matrices a (3*2 matrix of zeros) and b. I want to replace the 0s in matrix a with 1s depending on the index location stored in index b. I want to do it without a for loop.

import numpy as np

a = np.zeros((3, 2))
b = np.array([0, 1, 0])

The output should be

c = [[1, 0],
     [0, 1],
     [1, 0]]
indexing matrix numpy python
1 Answers

Numpy fancy indexing is your friend here. To make b work as the column index, you need an array of equal size to b which will tell you the row that each index applies to:

a[np.arange(b.size), b] = 1

This modifies a in-place. If that is not what you want, make a copy first:

c = a.copy()
c[np.arange(b.size), b] = 1
user_29887300
Related