python - numpy:将argsort应用于数组

argsort()函数返回一个索引矩阵，可用于索引原始数组，以便结果与sort()结果匹配。
有没有办法应用这些指数？我有两个数组，一个是用于获取排序顺序的数组，另一个是一些相关联的数据。
我想计算一下，但这似乎不起作用。
下面是一个例子：

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（或另一个具有相关数据的数组），但我不确定如何应用。

最佳答案

这可能是杀伤力过大，但在第二种情况下会起作用：

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]