让numpy的argsort返回一个二维索引数组？

Question

如果我们有一个一维数组。

arr = np.random.randint(7, size=(5))
# [3 1 4 6 2]
print np.argsort(arr)
# [1 4 0 2 3] <= The indices in the sorted order    

如果我们有一个二维数组。

arr = np.random.randint(7, size=(3, 3))
# [[5 2 4]
# [3 3 3]
# [6 1 2]]
print np.argsort(arr)
# [[1 2 0]
# [0 1 2]
# [1 2 0]] <= It sorts each row

我需要的是对这个矩阵进行完整排序后得到的二维索引。就像这样：

# [[2 1] => 1
# [0 1] => 2
# [2 2] => 2
# .
# .
# .
# [0 2] => 4
# [0 0] => 5
# [2 0]] => 6

如何为排序2d数组获取“2d索引”？

- AturSams

Answer 1

将 numpy.argsort 应用于扁平化数组，然后将索引展开回 (3, 3) 形状:

>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> np.dstack(np.unravel_index(np.argsort(arr.ravel()), (3, 3)))
array([[[2, 1],
        [0, 1],
        [2, 2],
        [1, 0],
        [1, 1],
        [1, 2],
        [0, 2],
        [0, 0],
        [2, 0]]])

- Ashwini Chaudhary

---

哎呀，这是什么结果啊...？ - MartianMartian

我认为你需要在末尾加上 [0]？ - Dirk Groeneveld

这看起来比它应该的要复杂。 - pietz

Answer 2

根据numpy.argsort的文档:

ind = np.unravel_index(np.argsort(x, axis=None), x.shape)

一个N维数组的排序元素的索引。

例如：

>>> x = np.array([[0, 3], [2, 2]])
>>> x
array([[0, 3],
       [2, 2]])
>>> ind = np.unravel_index(np.argsort(x, axis=None), x.shape)
>>> ind # a tuple of arrays containing the indexes
(array([0, 1, 1, 0]), array([0, 0, 1, 1]))
>>> x[ind]  # same as np.sort(x, axis=None)
array([0, 2, 2, 3])

- Mozak

---

为了从ind中删除与x中的np.nan值对应的索引，需要进行哪些编辑？ - Ouroboroski

Answer 3

使用np.take_along_axis来使用索引对切片进行操作。假设您需要按行排序：

>>> arr = np.array([[5, 2, 4],
[3, 3, 3],
[6, 1, 2]])
>>> ind = arr.argsort(axis=-1)
>>> ind
array([[1, 2, 0],
       [0, 1, 2],
       [1, 2, 0]])
>>> np.take_along_axis(arr, ind, axis=-1)
array([[2, 4, 5],
       [3, 3, 3],
       [1, 2, 6]])

- Ivelin

---

这并没有对所有元素进行排序，它只是对行进行了排序。 - AturSams