按元素排序多个numpy数组

Question

Suppose I have 3 arrays (I know the number of arrays beforehand) that I want to sort, element by element. For example, suppose I have:

import numpy as np

x = np.array([
[100, 200, 300],
[400, 500, 600],
[700, 800, 900]
])

y = np.array([
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
])

z = np.array([
[10, 20, 30],
[40, 50, 60],
[70, 80, 90]
])

And I'd want to get the order for each element:

[
[[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]],
[[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]],
[[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]], [[0, 2, 1], [0, 2, 1], [0, 2, 1]]
]

In this case all the positions are on the same order to keep the example simple (0 > 2 > 1), but in reality each element on each array could be bigger than its counterpart on another array. I'm also working with higher dimension arrays (3d rather than 2d), but I think the general idea would be the same.

I believe I need to reshape the arrays and then np.argsort them (or just regular np.sort , I'm more interested in the values rather than where they came from, though if the overhead of having the origin isn't high, it might be useful down the line). But I'm not sure how to reshape them such that I can apply the operation. All arrays are guaranteed to have the same shape, and despite my example here, it will always be an even number.

I'm not sure if the format I have in mind would be good for indexing after I have the order, but at the moment I can't figure out a better way to group the results. I believe I can just split them later with something like answer[...,0] for the maximum indexes (or maximum values), answer[...,1] for the second and so on. Though if I have the indexes, I'd have to map them later on to the arrays. Shouldn't be that hard I think, I'd just have to np.stack them together I believe.

Any suggestions?

Answer 1

You can use dstack() to bind them together over the last axis and then use np.sort() over the expected axis to sort the array out:

In [10]: arr = np.dstack((x, y, z))

In [11]: arr.sort(2)

In [12]: arr
Out[12]: 
array([[[  1,  10, 100],
        [  2,  20, 200],
        [  3,  30, 300]],

       [[  4,  40, 400],
        [  5,  50, 500],
        [  6,  60, 600]],

       [[  7,  70, 700],
        [  8,  80, 800],
        [  9,  90, 900]]])

And if you want the result in a descending order:

In [13]: arr[:,:,::-1]
Out[13]: 
array([[[100,  10,   1],
        [200,  20,   2],
        [300,  30,   3]],

       [[400,  40,   4],
        [500,  50,   5],
        [600,  60,   6]],

       [[700,  70,   7],
        [800,  80,   8],
        [900,  90,   9]]])