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

Question

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?

@unutbu Well sort of different here, as we are selecting one element per row as opposed to multiple elements per row in the linked previous question. — Divakar
– Divakar, Commented Jun 18, 2016 at 14:06
@Divakar: If the OP wants the 2D array, C, your answer on the linked page gives the desired result exactly. — unutbu
– unutbu, Commented Jun 18, 2016 at 14:15
@unutbu, the answer you linked use np.take, and I don't think it could fix my problem? — avocado
– avocado, Commented Jun 19, 2016 at 0:21

Divakar · Accepted Answer · 2016-11-18 13:00:22Z

7

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

edited Nov 18, 2016 at 13:00

answered Jun 18, 2016 at 13:56

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

avocado Over a year ago

so no broadcasting at all?

Divakar Over a year ago

Well you don't need broadcasting here, if I got the expected output correctly from the question. We are basically using integers at each dim, to select elements. We have the 2nd dim indices from B, so we just needed to create the corresponding ones for the first dim using np.arange. Hope that made sense!

Rahul · Accepted Answer · 2016-06-18 13:59:47Z

-1

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

answered Jun 18, 2016 at 13:59

Rahul

5128 silver badges16 bronze badges

Collectives™ on Stack Overflow

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

2 Answers 2

2 Comments

Comments

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Linked

Related