多次元ndarrayのargsort

Question

多次元配列を最後の軸で並べ替えるためのインデックスを取得しようとしています。

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

そして、私は、このようなnumpy.argsortの文書に基づいて

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

が、私はそれがこれを行う必要があると思ったが、私はエラーを取得しています、そのインデックスiたい：

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

を

---

編集：同じ形状の他の配列を再配置する必要があります（たとえば、配列bこのように同じ方法でa.shape == b.shape）が...私は、誰かが同じ問題を抱えて、the answer hereを見つけ

>>> 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に対してブロードキャストできるので、同様に機能します。

私はjのための '短い手' として

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

を考えます。それらは一緒に、新しい配列の各要素のソース行と列を定義します。彼らは一緒に、逐次ではなく動作します。

新しい配列へ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

ように。彼らキーはただ

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