python - Replacing value in a numpy array's column given index from a vector with the same rows

Question

A = np.zeros(shape = (3,4))
A = [[0 0 0 0]
     [0 0 0 0]
     [0 0 0 0]]

B = np.asarray[[2],[0],[3]]

Without a for loop, is there a simple way to change the value of the a components in A (to 1 for example), given the index in B

Such that:

A = [[0 0 1 0]
     [1 0 0 0]
     [0 0 0 1]]

I have been able to get this output with a for loop but would prefer it if it didn't for scalability in higher dimensional arrays.

Answer 1

Take a look at numpy's advanced indexing .

You can do

A[[0, 1, 2], [2, 0, 3]] = 1

We are indexing to fetch 3 locations and set a value, thus the lists used to index are each len 3. And A has 2 dims, so we use two lists (you can also use only one list if you want to fetch full rows).

It also works with arrays

A[np.arange(A.shape[0]), [2, 0, 3]] = 1