argsort для многомерного ndarray

Я пытаюсь получить индексы для сортировки многомерного массива по последней оси, например

>>> a = np.array([[3,1,2],[8,9,2]])

И я хотел бы индексы i такой, что

>>> a[i]
array([[1, 2, 3],
       [2, 8, 9]])

Основываясь на документации numpy.argsort, я думал, что это должно сделать это, но я получаю ошибку:

>>> a[np.argsort(a)]
IndexError: index 2 is out of bounds for axis 0 with size 2

Изменить: мне нужно переставить другие массивы той же формы (например, массив b такой, что a.shape == b.shape) так же... так что

>>> b = np.array([[0,5,4],[3,9,1]])
>>> b[i]
array([[5,4,0],
       [9,3,1]])
Источник

4 ответа

Решение

Решение:

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
       [2, 8, 9]])

Вы правильно поняли, хотя я бы не назвал это обманом индексации.

Может быть, это поможет прояснить ситуацию:

In [544]: i=np.argsort(a,axis=1)

In [545]: i
Out[545]: 
array([[1, 2, 0],
       [2, 0, 1]])

i это порядок, который мы хотим, для каждой строки. То есть:

In [546]: a[0, i[0,:]]
Out[546]: array([1, 2, 3])

In [547]: a[1, i[1,:]]
Out[547]: array([2, 8, 9])

Чтобы выполнить оба шага индексирования одновременно, мы должны использовать индекс "столбец" для 1-го измерения.

In [548]: a[[[0],[1]],i]
Out[548]: 
array([[1, 2, 3],
       [2, 8, 9]])

Еще один массив, который может быть в паре с i является:

In [560]: j=np.array([[0,0,0],[1,1,1]])

In [561]: j
Out[561]: 
array([[0, 0, 0],
       [1, 1, 1]])

In [562]: a[j,i]
Out[562]: 
array([[1, 2, 3],
       [2, 8, 9]])

Если i определяет столбец для каждого элемента, затем j определяет строку для каждого элемента. [[0],[1]] Массив столбцов работает так же хорошо, потому что он может быть передан против i,

Я думаю о

np.array([[0],
          [1]])

как "короткая рука" для j, Вместе они определяют исходную строку и столбец каждого элемента нового массива. Они работают вместе, а не последовательно.

Полное сопоставление с a к новому массиву есть:

[a[0,1]  a[0,2]  a[0,0]
 a[1,2]  a[1,0]  a[1,1]]

def foo(a):
    i = np.argsort(a, axis=1)
    return (np.arange(a.shape[0])[:,None], i)

In [61]: foo(a)
Out[61]: 
(array([[0],
        [1]]), array([[1, 2, 0],
        [2, 0, 1]], dtype=int32))
In [62]: a[foo(a)]
Out[62]: 
array([[1, 2, 3],
       [2, 8, 9]])
Источник

Приведенные выше ответы теперь немного устарели, поскольку в numpy 1.15 была добавлена ​​новая функциональность, чтобы упростить ее; take_along_axis ( https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.take_along_axis.html) позволяет сделать следующее:

>>> a = np.array([[3,1,2],[8,9,2]])
>>> np.take_along_axis(a, a.argsort(axis=-1), axis=-1)
array([[1 2 3]
       [2 8 9]])
Источник

Я нашел ответ здесь, с кем-то, имеющим ту же проблему. Они ключ просто обманывает индексирование, чтобы работать должным образом...

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
       [2, 8, 9]])
Источник

Вы также можете использовать linear indexing что может быть лучше с производительностью, вот так -

M,N = a.shape
out = b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]

Так, a.argsort(1)+(np.arange(M)[:,None]*N) в основном линейные индексы, которые используются для отображения b чтобы получить желаемый отсортированный вывод для b, Те же самые линейные индексы также могут быть использованы на a для получения отсортированного вывода для a,

Пробный прогон -

In [23]: a = np.array([[3,1,2],[8,9,2]])

In [24]: b = np.array([[0,5,4],[3,9,1]])

In [25]: M,N = a.shape

In [26]: b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]
Out[26]: 
array([[5, 4, 0],
       [1, 3, 9]])

Rumtime тесты -

In [27]: a = np.random.rand(1000,1000)

In [28]: b = np.random.rand(1000,1000)

In [29]: M,N = a.shape

In [30]: %timeit b[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
10 loops, best of 3: 133 ms per loop

In [31]: %timeit b.ravel()[a.argsort(1)+(np.arange(M)[:,None]*N)]
10 loops, best of 3: 96.7 ms per loop
Источник
Другие вопросы по тегам