配列の上位k点のインデックスを取得する

やりたいこと

題名の通り配列の上位k点のインデックスを取得する

##コード
ChatGptさんにお願いしたら書いてくれた

test.py

import heapq

# 配列
arr = [2, 5, 3, 7, 1, 8, 4]

# 上位k個のインデックスを取得する
k = 3

# indices = heapq.nlargest(k, range(len(arr)), arr.__getitem__)
#上の方の式が元々出力されたものだったが私にはarr.__getitem__が馴染みなかったのでlambda式で置き換えた
indices = heapq.nlargest(k, range(len(arr)), key = lambda x: arr[x])

# 結果を表示
print(indices)  # [5, 3, 1]

軽い説明

そもそものheapq.nlargestは

heapq.nlargest(n, iterable, key=None)
Return a list with the n largest elements from the dataset defined by iterable. key, if provided, specifies a function of one argument that is used to extract a comparison key from each element in the iterable: key=str.lower Equivalent to: sorted(iterable, key=key, reverse=True)[:n]

というふうに説明されていて
heapq.nlargest(k, data, key)
k: 取得したい値の個数
data: 対象
key: 比較するもの
となっている

今回は上位k点の値ではなくインデックスが欲しいので
data=配列のインデックスの範囲
key=それぞれのインデックスにおける配列の値
としている

他にもdic型であるkeyに対応するvalueで上位k点が欲しい時は

dict.py

import heapq

# データ
data = [    {"name": "Alice", "age": 20},    {"name": "Bob", "age": 25},    {"name": "Charlie", "age": 30},]

# 上位k個の要素を取得する
k = 2
largest = heapq.nlargest(k, data, key=lambda x: x["age"])

# 結果を表示
print(largest)  # [{"name": "Charlie", "age": 30}, {"name": "Bob", "age": 25}]

というようにかける

まとめ

ちょくちょく必要になるけど毎回忘れてしまう処理に関してまとめました
ChatGPTめちゃくちゃ便利なのでこれからも使っていきたいです