用另一个一维数组对ND numpy数组进行排序

Question

从这个问题的答案（使用一个较少的内存，沿着一个特定轴a对另一个numpy数组进行排序）中，我学习了如何根据另一个numpy数组的值对多维numpy数组进行排序，b而又不创建太多额外的数组。

但是，numpy.rec.fromarrays([a, b])仅当数组a和b具有相同形状时才有效。我的b数组是一维数组，但a数组是ND数组（未指定N）。a通过一维数组的值在特定轴上对数组进行排序是一种好方法（高效）b吗？

Answer 1

使用np.take与axis关键字参数：

>>> a = np.arange(2*3*4).reshape(2, 3, 4)
>>> a
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]]])
>>> b = np.arange(3)
>>> np.random.shuffle(b)
>>> b
array([1, 0, 2])
>>> np.take(a, b, axis=1)
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

如果要使用花式索引，则只需在索引元组中填充足够的空片即可：

>>> a[:, b]
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

或更一般的情况下：

>>> axis = 1
>>> idx = (slice(None),) * axis + (b,)
>>> a[idx]
array([[[ 4,  5,  6,  7],
        [ 0,  1,  2,  3],
        [ 8,  9, 10, 11]],

       [[16, 17, 18, 19],
        [12, 13, 14, 15],
        [20, 21, 22, 23]]])

但np.take实际上应该是您的首选。