如何使用 numpy.argsort() 作为二维以上的索引？

Posted 2023-02-23

【中文标题】如何使用 numpy.argsort() 作为二维以上的索引？

【英文标题】：How to use numpy.argsort() as indices in more than 2 dimensions?

【发布时间】：2018-06-12 19:27:24

【问题描述】：

我知道类似这个问题的问题已经被问过很多次了，但是对类似问题的所有答案似乎只适用于二维数组。

我对@9​​87654321@的理解是np.sort(array) == array[np.argsort(array)]应该是True。 我发现如果np.ndim(array) == 2 确实是正确的，但如果np.ndim(array) > 2 会给出不同的结果。

例子：

>>> array = np.array([[[ 0.81774634,  0.62078744],
                       [ 0.43912609,  0.29718462]],
                      [[ 0.1266578 ,  0.82282054],
                       [ 0.98180375,  0.79134389]]])
>>> np.sort(array)
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])
>>> array.argsort()
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]])
>>> array[array.argsort()]
array([[[[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]],



       [[[[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]],

         [[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]]])

那么，谁能向我解释一下np.argsort() 究竟是如何用作索引来获取排序数组的？ 我能想到的唯一方法是：

args = np.argsort(array)
array_sort = np.zeros_like(array)
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        array_sort[i, j] = array[i, j, args[i, j]]

这非常繁琐，无法针对任何给定数量的维度进行概括。

【参考方案1】：

这是一个通用的方法：

import numpy as np

array = np.array([[[ 0.81774634,  0.62078744],
                   [ 0.43912609,  0.29718462]],
                  [[ 0.1266578 ,  0.82282054],
                   [ 0.98180375,  0.79134389]]])

a = 1 # or 0 or 2

order = array.argsort(axis=a)

idx = np.ogrid[tuple(map(slice, array.shape))]
# if you don't need full ND generality: in 3D this can be written
# much more readable as
# m, n, k = array.shape
# idx = np.ogrid[:m, :n, :k]

idx[a] = order

print(np.all(array[idx] == np.sort(array, axis=a)))

输出：

True

说明：我们必须为输出数组的每个元素指定输入数组对应元素的完整索引。因此，输入数组的每个索引都具有与输出数组相同的形状，或者必须可广播到该形状。

我们不排序/argsort 的轴的索引保持不变。因此，我们需要为每一个传递一个可广播的范围（array.shape[i]）。最简单的方法是使用 ogrid 为所有维度创建这样一个范围（如果我们直接使用它，数组将返回原样。）然后将与排序轴对应的索引替换为 argsort 的输出。

2019 年 3 月更新：

Numpy 在强制多轴索引作为元组传递方面变得更加严格。目前，array[idx] 将触发弃用警告。为了成为未来的证明，请改用array[tuple(idx)]。 （感谢@Nathan）

或者使用numpy新的（1.15.0版）便利功能take_along_axis：

np.take_along_axis(array, order, a)

【讨论】：

无论如何，+1 表示通用。

经过一番摆弄，我成功地将您的解决方案实现到我的代码中。 +1 完全回答了我的问题并提供了非常有用的解决方案。

您现在需要添加idx = tuple(idx) 以避免FutureWarning（以及后来的错误/错误结果）。 argsort 文档中确实需要这样的东西；应该有一种简单/标准的方法可以从排序索引到排序数组。

@Nathan Done，我也部分解决了您的最后一点。也许take/put_along_axis 应该从arg* 函数文档链接。

感谢您的快速更新和添加take_along_axis。我注意到take_along_axis（以及put_along_axis）的文档确实引用了argsort 和argpartition，所以看起来这些确实是现在应该完成的标准方式；如果有来自argsort/argprartition 文档的另一个方向的参考，那确实很棒。

【参考方案2】：

@Hameer 的回答有效，尽管它可能会使用一些简化和解释。

sort 和 argsort 正在最后一个轴上工作。 argsort 返回一个 3d 数组，形状与原始数组相同。这些值是最后一个轴上的索引。

In [17]: np.argsort(arr, axis=2)
Out[17]: 
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]], dtype=int32)
In [18]: _.shape
Out[18]: (2, 2, 2)
In [19]: idx=np.argsort(arr, axis=2)

要使用它，我们需要为广播到相同 (2,2,2) 形状的其他维度构建索引。 ix_ 是一个方便的工具。

仅使用idx 作为ix_ 输入之一是行不通的：

In [20]: np.ix_(range(2),range(2),idx)
....
ValueError: Cross index must be 1 dimensional

我使用最后一个范围，然后忽略它。 @Hameer 改为构造 2d ix_，然后展开它们。

In [21]: I,J,K=np.ix_(range(2),range(2),range(2))
In [22]: arr[I,J,idx]
Out[22]: 
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])

因此其他维度的索引与(2,2,2) idx 数组一起使用：

In [24]: I.shape
Out[24]: (2, 1, 1)
In [25]: J.shape
Out[25]: (1, 2, 1)

当您获得一维的多维索引时，这是构造其他索引的基础。

@Paul 构造与ogrid 相同的索引：

In [26]: np.ogrid[slice(2),slice(2),slice(2)]  # np.ogrid[:2,:2,:2]
Out[26]: 
[array([[[0]],

        [[1]]]), array([[[0],
         [1]]]), array([[[0, 1]]])]
In [27]: _[0].shape
Out[27]: (2, 1, 1)

ogrid 作为class 使用切片，而ix_ 需要列表/数组/范围。

argsort for a multidimensional ndarray（从 2015 年开始）适用于二维数组，但同样的逻辑也适用（查找使用 argsort 广播的范围索引）。

【参考方案3】：

这是一个矢量化的实现。它应该是 N 维的，并且比你正在做的要快很多。

import numpy as np


def sort1(array, args):
    array_sort = np.zeros_like(array)
    for i in range(array.shape[0]):
        for j in range(array.shape[1]):
            array_sort[i, j] = array[i, j, args[i, j]]

    return array_sort


def sort2(array, args):
    shape = array.shape
    idx = np.ix_(*tuple(np.arange(l) for l in shape[:-1]))
    idx = tuple(ar[..., None] for ar in idx)
    array_sorted = array[idx + (args,)]

    return array_sorted


if __name__ == '__main__':
    array = np.random.rand(5, 6, 7)
    idx = np.argsort(array)

    result1 = sort1(array, idx)
    result2 = sort2(array, idx)

    print(np.array_equal(result1, result2))