Ранг 2D массив по строкам

Question

Этот вопрос задает вопрос о сортировке одномерных массивов, даны два решения, первое предлагает с использованием argsort дважды , второе, более эффективное по времени, использует его только один раз . Что если бы я хотел ранжировать двумерный массив по очереди, как это?

Использование 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+ с предложенным для такого массива.