对 numpy 数组进行排名以返回排名数组

Question

我有一个 3 维 NumPy 数组，表示 n 个维度为 x,x 的矩阵。例如，当 n=3 和 x=2 时，示例如下：

matrices = np.array([[[2,8],[1,7]],[[3,1],[5,4]],[[9,6],[2,3]]]

matrices
array([[[2, 8],
        [1, 7]],

       [[3, 1],
        [5, 4]],

       [[9, 6],
        [2, 3]]])

我想创建一个形状相同的新数组，但值是每个元素的这些矩阵的排名（即沿轴 0 排名）。结果如下：

array([[[1, 3],
        [1, 3]],

       [[2, 1],
        [3, 2]],

       [[3, 2],
        [2, 1]]])

我可以看到如何排序 (np.argsort(matrices, axis=0))，但找不到返回排名的简单方法。矩阵尺寸 x 可能很高，因此快速运行时间很重要。使用的python的分布也不包括scipy。

我发现以下工作：

def rank_stations(input):
    output = [0] * (np.size(input))
    for i, x in enumerate(sorted(range(len(input)), reverse=True, key=lambda y: input[y])):
        output[x] = i+1
    return output

results = np.apply_along_axis(rank_stations, 0, matrices)

使用 5 个尺寸为 980 x 980 的矩阵，运行时间约为 6 秒，而建模软件计算为 20 秒，因此这是一个合理的改进。排序算法在几分之一秒内运行，那么有没有办法通过排名获得相似的运行时间？

Answer 1

将numpy.apply_along_axis 与scipy.stats.rankdata 一起使用：

from scipy.stats import rankdata

np.apply_along_axis(rankdata, 0, matrices)

或者使用np.argsort 两次：

f = lambda x: x.argsort().argsort()
np.apply_along_axis(f, 0, matrices) + 1

输出：

array([[[1., 3.],
        [1., 3.]],

       [[2., 1.],
        [3., 2.]],

       [[3., 2.],
        [2., 1.]]])