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