用一个小的N-D数组按列对N-D numpy数组进行排序

Question

我知道通过使用高级索引对N-D numpy数组按另一个1-D数组排序，可以使用fancy indexing实现c = a[:, :, b]，其中b定义了按列排序的顺序。

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

       [[ 5,  4],
        [ 7,  6]],

       [[ 9,  8],
        [11, 10]]])

现在我通过添加两个输入到b2来增加b，对应于我想要对a中每个2x2集合进行排序的方式。

>>> b2 = np.array([[1, 0], [0, 1], [1, 0]])
>>> c2 = ?
>>> c2
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

我有一组更大的输入数据，我有一个类似于“b2”的返回数组的函数，它提供了我应该获取哪些信息。因此，我想知道我应该填入c2 = ?以获得所需的结果。

- kent

---

对于期望输出，b2不应该是np.array([[1, 0], [0, 1], [1, 0]])（最后一行交换了位置）吗？ - Alex Riley

@ajcr 感谢您发现了那个错误！ :) - kent

Answer 1

这里有一种运用高级索引的方法-fancy-indexing-

(a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)

示例运行 -

In [191]: a
Out[191]: 
array([[[7, 8, 5, 2, 0],
        [6, 7, 0, 7, 1],
        [7, 6, 5, 4, 0]],

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

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

In [192]: b2
Out[192]: 
array([[1, 2, 4, 3, 0],
       [4, 2, 0, 1, 3],
       [1, 3, 4, 0, 2]])

In [193]: (a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)
Out[193]: 
array([[[8, 5, 0, 2, 7],
        [7, 0, 1, 7, 6],
        [6, 5, 0, 4, 7]],

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

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

- Divakar

---

这是一个很好的例子，展示了混合使用高级和基本索引的方法。 - unutbu

@unutbu，我甚至不知道我在混淆事情！ :) - Divakar

Answer 2

如果没有人找到一个纯粹的花式索引解决方案，这里有一个循环遍历第一个轴的解决方案:

np.asarray([a[n,:,p] for n,p in enumerate(b2)])

array([[[ 1,  3],
    [ 0,  2]],

   [[ 4,  6],
    [ 5,  7]],

   [[ 9, 11],
    [ 8, 10]]])

- Christian K.

---

谢谢你的回答！使用高级索引解决方案会更好，因为我正在寻找速度。 :) - kent

Answer 3

这与@Divakar的解决方案类似，但没有转置。

In [259]: I,J, K = np.ogrid[:3,:2,:2]

In [260]: a[I, J, b[:,None,:]]
Out[260]: 
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

我正在使用ogrid（或np.ix_）作为一种紧凑的方式来生成2个3D数组，这些数组与b [:，None，：]广播以产生一组（3,2,2）索引。

完全等效的写法是：Nones。

a[np.arange(3)[:,None,None], np.arange(2)[None,:,None], b[:,None,:]]

这样更清晰地表明了b是沿着第一个和最后一个轴选择项目。

要查看完整的广播索引数组，请打印：

 np.broadcast_arrays(I,J,b[:,None,:])

- hpaulj

---

谢谢！和Divakar的回答类似，我学到了很多，也感谢你逐步解释。 - kent