按特定轴通过另一个数组对numpy数组进行排序

Question

就像这个回答里说的，我有一对三维的numpy数组，a和b，我想根据a的值来给b的内容排序。不过和这个回答不一样，我只想在数组的一个方向上进行排序。

我对numpy.argsort()文档的简单理解是：

Returns
-------
index_array : ndarray, int
    Array of indices that sort `a` along the specified axis.
    In other words, ``a[index_array]`` yields a sorted `a`.

这让我觉得可以用下面的代码来进行排序：

import numpy

a = numpy.zeros((3, 3, 3))
a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
print "a"
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]

 [[ 3.  3.  3.]
  [ 3.  3.  3.]
  [ 3.  3.  3.]]

 [[ 2.  2.  2.]
  [ 2.  2.  2.]
  [ 2.  2.  2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]]
"""
print "a, sorted"
print numpy.sort(a, axis=0)
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]

 [[ 2.  2.  2.]
  [ 2.  2.  2.]
  [ 2.  2.  2.]]

 [[ 3.  3.  3.]
  [ 3.  3.  3.]
  [ 3.  3.  3.]]]
"""

##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:

[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""

那么，正确的做法是什么呢？

Answer 1

numpy.take_along_axis() 这个函数可以做到这一点，而且应该不会浪费额外的内存：

# From the original question:
import numpy
a = numpy.zeros((3, 3, 3))
a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
b = numpy.arange(3*3*3).reshape((3, 3, 3))
sort_indices = numpy.argsort(a, axis=0)
# This is not working as expected:
c = b[sort_indices]
# This does what is expected:
c = numpy.take_along_axis(b, sort_indices, axis=0)
print(c)
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]
"""

Answer 2

要让这个功能正常工作，你还需要为另外两个维度提供索引。

>>> a = numpy.zeros((3, 3, 3))
>>> a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
>>> b = numpy.arange(3*3*3).reshape((3, 3, 3))
>>> sort_indices = numpy.argsort(a, axis=0)
>>> static_indices = numpy.indices((3, 3, 3))
>>> b[sort_indices, static_indices[1], static_indices[2]]
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]]])

numpy.indices 这个函数可以计算数组每个轴的索引，前提是通过其他两个轴（或者说 n - 1 个轴，其中 n 是总的轴数）来“展平”数组。换句话说，就是这个（抱歉内容有点长）：

>>> static_indices
array([[[[0, 0, 0],
         [0, 0, 0],
         [0, 0, 0]],

        [[1, 1, 1],
         [1, 1, 1],
         [1, 1, 1]],

        [[2, 2, 2],
         [2, 2, 2],
         [2, 2, 2]]],


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

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

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


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

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

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

这些是每个轴的身份索引；当用来索引 b 时，它们可以重新创建 b。

>>> b[static_indices[0], static_indices[1], static_indices[2]]
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

除了 numpy.indices，你还可以使用 numpy.ogrid，正如 unutbu 所建议的那样。由于 ogrid 生成的对象更小，我会创建所有三个轴，仅仅是为了保持一致，但请注意 unutbu 的评论，里面有提到只生成两个轴的方法。

>>> static_indices = numpy.ogrid[0:a.shape[0], 0:a.shape[1], 0:a.shape[2]]
>>> a[sort_indices, static_indices[1], static_indices[2]]
array([[[ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.]],

       [[ 2.,  2.,  2.],
        [ 2.,  2.,  2.],
        [ 2.,  2.,  2.]],

       [[ 3.,  3.,  3.],
        [ 3.,  3.,  3.],
        [ 3.,  3.,  3.]]])