Question

This question询问有关对一维数组进行排序的问题，给出了两种解决方案，第一种建议使用using argsort twice，第二种建议使用更省时的uses it only once。如果我想像这样按行排列2D数组怎么办？

两次使用argsort是一种可能：

def rank(x, axis=-1):
   return np.argsort(np.argsort(x, axis=axis), axis=axis)

有关数据：

x = np.array([
    [1,  2,  30],
    [4,  5,  6],
    [90, 8,  7],
    [12, 15, 10]
])

它返回正确的结果：

rank(x, axis=0)
## array([[0, 0, 3],
##        [1, 1, 0],
##        [3, 2, 1],
##        [2, 3, 2]])

rank(x, axis=1)
## array([[0, 1, 2],
##        [0, 1, 2],
##        [2, 1, 0],
##        [1, 2, 0]])

但是有没有更有效的方法？

Answer 1

这是更好的链接解决方案的2D版本。

def rank(x, axis=-1):
    m, n = x.shape
    res = np.empty((m, n), dtype=int)
    I = np.ogrid[:m, :n]
    rng, I[axis] = I[axis], x.argsort(axis=axis)
    res[I] = rng
    return res

和ND版本：

def rank(x, axis=-1):
    res = np.empty(x.shape, dtype=int)
    I = np.ogrid[tuple(map(slice, x.shape))]
    rng, I[axis] = I[axis], x.argsort(axis=axis)
    res[I] = rng
    return res

Answer 2

Using only once与advanced-indexing-

sidx = np.argsort(x, axis=1)

# Store shape info
m,n = x.shape

# Initialize output array
out = np.empty((m,n),dtype=int)

# Use sidx as column indices, while a range array for the row indices
# to select one element per row. Since sidx is a 2D array of indices
# we need to use a 2D extended range array for the row indices
out[np.arange(m)[:,None], sidx] = np.arange(n)

要获得各列的排名，只需更改索引步骤以使用sidx作为行索引，而列索引的范围数组将自动广播。同样，要分配的值将扩展为2D，以便在分配之前广播这些值：

sidx = np.argsort(x, axis=0)
out[sidx, np.arange(n)] = np.arange(m)[:,None]

5k x 5k数组上的时间

要查看所提出方法的改进，因为两者都使用了第一个argsort的结果，我们先对其进行预先计算，然后安排其余步骤的时间-

In [248]: x = np.random.rand(5000,5000)

In [249]: axis = 1

In [250]: sidx = np.argsort(x, axis=1)

In [251]: %timeit np.argsort(sidx, axis=axis)
1 loop, best of 3: 1.31 s per loop

In [252]: %%timeit
     ...: m,n = x.shape
     ...: out = np.empty((m,n),dtype=int)
     ...: out[np.arange(m)[:,None], sidx] = np.arange(n)
10 loops, best of 3: 156 ms per loop

对于这样的数组，建议的8x+加速了。

逐行排列2D数组

2 个答案: