Sort numpy matrix based on index-matrix

Question

In numpy, we can sort arrays like this:

>>> import numpy as np

>>> a = np.array([0, 100, 200])
>>> order = np.array([1, 2, 0])
>>> print(a[order])
[100 200   0]

However, this does not work when the "order" is a matrix:

>>> A = np.array([    [0, 1, 2],
                      [3, 4, 5],
                      [6, 7, 8]])

>>> Ord = np.array([  [1, 0, 2],
                      [0, 2, 1],
                      [2, 1, 0]])

>>> print(A[Ord].shape)
(3, 3, 3)

I would like to have "A" sorted like this:

array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

Kevin · Accepted Answer · 2021-04-02 15:14:15Z

2

You could use np.take_along_axis for this.

np.take_along_axis(A, Ord, axis=1)

Output

array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

As stated in the documentation it is often used together with functions that produce indices, like argsort. But I am not sure if this would generalize for more than 2 dimensions.

answered Apr 2, 2021 at 15:14

Kevin

3,4182 gold badges16 silver badges55 bronze badges

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

Comments

hasref · Accepted Answer · 2021-04-02 10:55:43Z

There are two general ways of indexing into a numpy array:

Basic indexing and slicing which extends python's concept of slicing.
Advanced Indexing where we generally use another ndarray containing integers/booleans.

You can read about these in the numpy documentation.

To answer your question: we can make use of advanced indexing here. The documentation contains an example, adapted below, which is already quite close to what we want:

>>> import numpy as np
>>> A = np.array([[1,2],
                  [3,4]])
>>> row_indices = np.array([[0,0],
                            [1,1]])
>>> col_indices = np.array([[1,0],
                            [1,0]])
>>> A[row_indices,col_indices]
array([[2, 1],
       [4, 3]])

In the code in the question, Ord already contains the column indices, so all we need to do is generate the row indices ourselves. While this is probably not the best way to do this, here is one possible solution:

>>> A = np.array( [[0,1,2],
...                [3,4,5],
...                [6,7,8]])
>>> col_indices = np.array([[1,0,2],
...                         [0,2,1],
...                         [2,1,0]])
>>> row_indices = np.repeat([0,1,2],3).reshape(3,3)
>>> row_indices
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2]])
>>> A[row_indices, col_indices]
array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

Collectives™ on Stack Overflow

Sort numpy matrix based on index-matrix

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related