CODEWITHSUNDEEP
pythonarraysnumpy
Alcott
Alcott

Reputation: 18605

Select elements from an array using another array as index

Say I have an array

A = array([[1,2,3],
           [4,5,6],
           [7,8,9]])

Index array is 

B = array([[1], # want [0, 1] element of A
           [0],  # want [1, 0], element of A
           [1]])  # want [2, 1] elemtn of A

By this index array B, I want a 3-by-1 array, whose elements are taken from array A, that is 

C = array([[2],
           [4],
           [8]])

I tried numpy.choose, but I failed to that.

Upvotes: 4

Views: 1535

Answers (3)

Jaime
Jaime

Reputation: 67507

For answer completeness... Fancy indexing arrays are broadcast to a common shape, so the following also works, and spares you that final reshape:

>>> A[np.arange(3)[:, None], B]
array([[2],
       [4],
       [8]])

Upvotes: 5

CT Zhu
CT Zhu

Reputation: 54380

You can do:

>>>np.diag(A[range(3),B]).reshape(B.shape)
array([[2],
       [4],
       [8]])

If you want to use choose you can do: np.choose(B.ravel(), A.T).reshape(B.shape).

Upvotes: 2

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 251156

You can do something like this:

>>> A[np.arange(len(A)), B.ravel()].reshape(B.shape)
array([[2],
       [4],
       [8]])

Upvotes: 3

Related Questions