numpy：将argsort应用于数组

Question

argsort()函数返回一个索引矩阵，可以用来索引原始数组以便结果与sort()函数的结果相匹配。

是否有一种方法可以应用这些索引？我有两个数组，一个用于获取排序顺序的数组，另一个是一些相关数据。

我想计算assoc_data[array1.argsort()]，但似乎不起作用。

以下是一个示例：

z=array([1,2,3,4,5,6,7])
z2=array([z,z*z-7])
i=z2.argsort()

---

z2=array([[ 1,  2,  3,  4,  5,  6,  7],
          [-6, -3,  2,  9, 18, 29, 42]])
i =array([[1, 1, 1, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 1, 1]])

我想将i应用于z2（或另一个带有相关数据的数组），但我不确定如何做到这一点。

- Jason S

---

你想沿着哪个轴排序？ - user545424

Answer 1

这可能有些过度，但这将在第二种情况下起作用：

import numpy as np
axis = 0
index = list(np.ix_(*[np.arange(i) for i in z2.shape]))
index[axis] = z2.argsort(axis)
z2[index]

# Or if you only need the 3d case you can use np.ogrid.

axis = 0
index = np.ogrid[:z2.shape[0], :z2.shape[1], :z2.shape[2]]
index[axis] = z2.argsort(axis)
z2[index]

- Bi Rico

---

很奇怪，我一年后再次需要这个，而在寻找如何做到这一点时，我偶然发现了我之前问过的问题... 我终于明白这是干什么用的了。顺便说一句，它看起来并不过度杀伤。 - Jason S

Answer 2

你很幸运，我刚刚获得了numpy学位。

>>> def apply_argsort(a, axis=-1):
...     i = list(np.ogrid[[slice(x) for x in a.shape]])
...     i[axis] = a.argsort(axis)
...     return a[i]
... 
>>> a = np.array([[1,2,3,4,5,6,7],[-6,-3,2,9,18,29,42]])
>>> apply_argsort(a,0)
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

如果要了解正在发生的事情，请查看我的回答这篇问题。

- user545424

Answer 3

Use np.take_along_axis

np.take_along_axis(z2, i, axis=1)
Out[31]: 
array([[ 1,  2,  3,  4,  5,  6,  7],
       [-6, -3,  2,  9, 18, 29, 42]])

- vozman

---

这不是我要找的答案，但如果正确调用，np.take_along.axis 确实可以用于此目的。 - Jason S

Answer 4

啊哈，搞定了。

In [274]: z2[i,range(z2.shape[1])]
Out[274]:
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

- Jason S

---

你还没说我是什么 ;) - tiho

抱歉，我被你的问题中有两个“i”的事实搞混了，我看错了“错误”的那一个（第一个），当它插入你的答案时不起作用。 - tiho