将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)

[使用五个尺寸为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.]]])