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

Question

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

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

Answer 1

Решение:

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

Answer 2

Я нашел the answer here, с кем-то с той же проблемой. Они ключ просто обманывают индексации, чтобы работать должным образом ...

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

Answer 3

Вы также можете использовать 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