python - 用另一个一维数组对 N 维 numpy 数组进行排序

Question

从这个问题 (Sort a numpy array by another array, along a particular axis, using less memory) 的答案中，我学会了如何根据另一个 numpy 数组 b 的值对多维 numpy 数组 a 进行排序，而无需创建太许多额外的数组。

但是，numpy.rec.fromarrays([a, b]) 仅在数组 a 和 b 具有相同形状时才有效.我的 b 数组是一维数组，但 a 数组是 N 维数组(未指定 N)。根据一维数组 b 的值对特定轴中的 a 数组进行排序是否是一种很好的方法(并且有效)？

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 应该是您的第一选择。