numpyの多次元配列からk個の最大値のインデックスを取得する方法

Question

StackOverflowに関するいくつかの質問がありましたが、関連する回答が見つかりませんでした。私はnumpyndarrayからk個の最大値のインデックスをフェッチしたいと思います。 This linkは1Dアレイでも同じですが、説明しています。 2D配列用のnp.argsortは、要素を行ごとにソートしました。すなわち

Note: array elements are not unique. 

入力：

import numpy as np 
n = np.arange(9).reshape(3,3) 
>>> n 
array([[0, 1, 2], 
    [3, 4, 5], 
    [6, 7, 8]]) 
s = n.argsort() 
>>> s 
array([[0, 1, 2], 
    [0, 1, 2], 
    [0, 1, 2]], dtype=int32) 

また、

import numpy as np 
n = np.arange(9).reshape(3,3) 
s = n.argsort(axis=None) 
>>>s 
array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=int32) 

が、私はここに配列構造を失うていますし、要素の元のインデックスを償還することはできません。

何か助けがありがたいです。 ndarraysためnp.argpartitionとnp.argsortとアプローチの

Answer 1

カップル - k_largest_index_argpartition_v1とk_largest_index_argpartition_v2間argpartition

異なる2つのバージョンの

def k_largest_index_argpartition_v1(a, k): 
    idx = np.argpartition(-a.ravel(),k)[:k] 
    return np.column_stack(np.unravel_index(idx, a.shape)) 

def k_largest_index_argpartition_v2(a, k): 
    idx = np.argpartition(a.ravel(),a.size-k)[-k:] 
    return np.column_stack(np.unravel_index(idx, a.shape)) 

def k_largest_index_argsort(a, k): 
    idx = np.argsort(a.ravel())[:-k-1:-1] 
    return np.column_stack(np.unravel_index(idx, a.shape)) 

考察は、我々はargparitionを使用する方法です。最初のバージョンでは、入力配列を無効にしてから、kインデックスのインデックスを取得するためにargpartitionを使用しているため、最大のkインデックスが効果的に取得されますが、2番目のバージョンでは最小インデックスはa.size-k、残りの最大のkインデックスを選択します。

また、argpartitionでは、ソートされた順序でインデックスが取得されません。ソートされた順序が必要な場合は、postで説明したように、範囲配列をnp.argpartitionにフィードする必要があります。

サンプルの実行 -

1）2次元の場合：

In [42]: a # 2D array 
Out[42]: 
array([[38, 14, 81, 50], 
     [17, 65, 60, 24], 
     [64, 73, 25, 95]]) 

In [43]: k_largest_index_argsort(a, k=2) 
Out[43]: 
array([[2, 3], 
     [0, 2]]) 

In [44]: k_largest_index_argsort(a, k=4) 
Out[44]: 
array([[2, 3], 
     [0, 2], 
     [2, 1], 
     [1, 1]]) 

In [66]: k_largest_index_argpartition_v1(a, k=4) 
Out[66]: 
array([[2, 1], # Notice the order is different 
     [2, 3], 
     [0, 2], 
     [1, 1]]) 

2）3Dの場合：

In [46]: a # 3D array 
Out[46]: 
array([[[20, 98, 27, 73], 
     [33, 78, 48, 59], 
     [28, 91, 64, 70]], 

     [[47, 34, 51, 19], 
     [73, 38, 63, 94], 
     [95, 25, 93, 64]]]) 

In [47]: k_largest_index_argsort(a, k=2) 
Out[47]: 
array([[0, 0, 1], 
     [1, 2, 0]]) 

ランタイムテスト -

In [56]: a = np.random.randint(0,99999999999999,(3000,4000)) 

In [57]: %timeit k_largest_index_argsort(a, k=10) 
1 loops, best of 3: 2.18 s per loop 

In [58]: %timeit k_largest_index_argpartition_v1(a, k=10) 
10 loops, best of 3: 178 ms per loop 

In [59]: %timeit k_largest_index_argpartition_v2(a, k=10) 
10 loops, best of 3: 128 ms per loop