using an numpy array as indices of the 2nd dim of another array?

For example, I have two numpy arrays,

A = np.array(
  [[0,1], 
   [2,3], 
   [4,5]])
B = np.array(
  [[1],
   [0],
   [1]], dtype='int')

and I want to extract one element from each row of A, and that element is indexed by B, so I want the following results:

C = np.array(
  [[1],
   [2],
   [5]])

I tried A[:, B.ravel()], but it’ll broadcast B, not what I want. Also looked into np.take, seems not the right solution to my problem.

However, I could use np.choose by transposing A,

np.choose(B.ravel(), A.T)

but any other better solution?

Answers:

Thank you for visiting the Q&A section on Magenaut. Please note that all the answers may not help you solve the issue immediately. So please treat them as advisements. If you found the post helpful (or not), leave a comment & I’ll get back to you as soon as possible.

Method 1

You can use NumPy's purely integer array indexing –

A[np.arange(A.shape[0]),B.ravel()]

Sample run –

In [57]: A
Out[57]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [58]: B
Out[58]: 
array([[1],
       [0],
       [1]])

In [59]: A[np.arange(A.shape[0]),B.ravel()]
Out[59]: array([1, 2, 5])

Please note that if B is a 1D array or a list of such column indices, you could simply skip the flattening operation with .ravel().

Sample run –

In [186]: A
Out[186]: 
array([[0, 1],
       [2, 3],
       [4, 5]])

In [187]: B
Out[187]: [1, 0, 1]

In [188]: A[np.arange(A.shape[0]),B]
Out[188]: array([1, 2, 5])

Method 2

C = np.array([A[i][j] for i,j in enumerate(B)])

All methods was sourced from stackoverflow.com or stackexchange.com, is licensed under cc by-sa 2.5, cc by-sa 3.0 and cc by-sa 4.0