How to access 2D numpy array elements along axis 1 using an array of indices?

Let N be a numpy array of shape (m,n) and A be an array of indices (m,) containing values between 0 and n-1 including.

If I i is an index in range from 0 to m-1 including, then I want to access elements N[i,A[i]].

For example:

import numpy as np

m = 2000
n = 10

N = np.zeros((m,n))
A = np.random.choice(n, m)

N[???, A] = 1

I expect the example code above to generate an array N of shape (m,n) initially with all zeros, and then set N[i,j]=1 where A[i]=j.

I tried N[:, A]=1, but instead it changes all the elements of the array N to 1, which is not what I’m trying to achieve.

>Solution :

Using : indexes all the rows, so doing N[:,A] = 1 is telling numpy to take all the rows and in column A set the value to 1. You want to just have a one-to-one row and column selection. The simplest way to do that is to use np.arange.

N[np.arange(m), A] = 1

Now, each value in A has a corresponding row, which results in m row and column selections.