numpy.partitionとnumpy.argpartitionの両方を効率的に出力する

Question

私はpython 3.6とnumpyを使用しています。私はn次元の配列を持っています。私は、配列の最後の次元でパーティションとargパーティションの両方を実行する必要があります。私は明らかに両方の関数を呼び出すことができますが、リソースを無駄にするような感じです。 np.partitionとnp.argpartitionの両方の結果を同時に取得する方法はありますか？ np.partitionの結果を、np.argpartitionから取得したインデックスを配列に適用する方法が必要ですが、現時点では表示されません！

ありがとうございました！

Answer 1

インデックスを取得し、advanced-indexingを使用して、分割された配列を取得します。したがって

、寸法の任意の数の一般的ndarrayための任意の一般的な軸に沿った実装、そうのようになる -

def partition_results(a, k, axis=-1): 
    idx = np.argpartition(a, k, axis=axis) 
    index_arr = list(np.ix_(*[range(i) for i in a.shape])) 
    index_arr[axis] = idx 
    return idx, a[index_arr] 

np.ix_は、タスクを達成するために私達に「広がるアウト」レンジアレイを与えますadvanced-indexing。これらの範囲配列は、argpartitionインデックス配列内の軸の長さに対応するすべての次元をカバーするために必要です。最後のインデックス配列を除くインデックス配列はargpartitionインデックスです。この設定は、このような索引付け操作に必要です。

そうは、np.argpartitionとnp.partitionへの2つの別々の呼び出しを使用してのアプローチでは、我々はそれを持っているでしょう、そうのように -

def partition_results_exclusive_way(a, k): 
    idx = np.argpartition(a, k, axis=-1) 
    part_arr = np.partition(a, k, axis=-1) 
    return idx , part_arr 

私たちは、次のセクションの性能と価値検証の上、比較のためにそれを使用します。

サンプル実行およびランタイムテスト -

In [496]: a = np.random.rand(20,20,20,20,20) 

In [502]: A0, B0 = partition_results_exclusive_way(a, 10) 

In [503]: A1, B1 = partition_results(a, 10) 

In [504]: np.allclose(A0,A1) 
Out[504]: True 

In [505]: np.allclose(B0,B1) 
Out[505]: True 

In [506]: %timeit partition_results_exclusive_way(a, 10) 
10 loops, best of 3: 92.6 ms per loop 

In [507]: %timeit partition_results(a, 10) 
10 loops, best of 3: 76 ms per loop 

パフォーマンス番号にもう少し解剖、レッツ・時間別にargpartitionとpartition - だから、

In [509]: %timeit np.argpartition(a, 10, axis=-1) 
10 loops, best of 3: 49.6 ms per loop 

In [510]: %timeit np.partition(a, 10, axis=-1) 
10 loops, best of 3: 43.6 ms per loop 

、advanced-indexing操作はの半分の周りに私たちを負けまし私たちが持っていたのはnp.partitionです。私たちは間違いなくそこに救います！