Sort numpy matrix based on index-matrix

user_6319376 0
Viewed 495

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]])
matrix numpy sorting arrays python
2 Answers

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.

user_149232270

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

  1. Basic indexing and slicing which extends python's concept of slicing.
  2. 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]])
user_51029110
Related