loganecolss loganecolss - 20 days ago 5
Python Question

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?

Answer

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])