Pythonでソートが使えない場合の最大値からX番目に大きい値までのインデックスの取得方法

概要

あまり需要はないと思いますが，Qiitaへの投稿に慣れるために，pythonで上位から任意の順位までの値のインデックスをソートを使わずにリストとして手に入れる方法を実装してみます．もっと効率的な方法があれば使いたい．

環境

numpy==1.14.0
Python 3.5.0

仮定

• 同じ数が存在しない（もしくは，先頭に近い方を取得すれば良い場合）

実装

Aが最大値のインデックスを手に入れたいデータ．TOPが何番目の値までかを示す．

A = [4,9,1]
TOP = 2
results = []
for num in range(TOP):
    idx = np.argmax(A)
    results.append(idx)
    A[idx] = float('-inf')

print(results)

結果はこんな感じ

$ # Aの最大値と2番目の最大値のインデックスが手に入る
$ [1, 0]

Aがarrayだと次の感じ（各行ごとで計算をしていると仮定しています，列ならばaxisの値を0にすればいいはず）

A = [[3,2,1],[4,5,6],[7,9,8]]
TOP = 2
results = []
for num in range(TOP):
    idxs = np.argmax(A,axis=1)
    results.append(idxs)
    for row, idx in zip(A, idxs):
        row[idx] = float('-inf')

print(results)

結果はこんな感じ

$ # 最初のarrayが最大値のインデックス，次のarrayが2番目の値の
$ # インデックスが手に入る
$ [array([0, 2, 1]), array([1, 1, 2])]