niekas
niekas

Reputation: 9087

Sort numpy.array rows by indices

I have 2D numpy.array and a tuple of indices:

a = array([[0, 0], [0, 1], [1, 0], [1, 1]])
ix = (2, 0, 3, 1)

How can I sort array's rows by the indices? Expected result:

array([[1, 0], [0, 0], [1, 1], [0, 1]])

I tried using numpy.take, but it works as I expect only with 1D arrays.

Upvotes: 3

Views: 5726

Answers (1)

NPE
NPE

Reputation: 500257

You can in fact use ndarray.take() for this. The trick is to supply the second argument (axis):

>>> a.take(ix, 0)
array([[1, 0],
       [0, 0],
       [1, 1],
       [0, 1]])

(Without axis, the array is flattened before elements are taken.)

Alternatively:

>>> a[ix, ...]
array([[1, 0],
       [0, 0],
       [1, 1],
       [0, 1]])

Upvotes: 6

Related Questions