yiskw note

機械学習やプログラミングについて気まぐれで書きます

【Python】リストから大きい/小さい順にn個の要素を取り出す


Pythonの標準ライブラリであるheapq.nlargest / heapq.nsmallestを使用することで,
リストから簡単に大きい/小さい順にn個の要素を取り出すことができます.

import heapq
import random


a = [random.randint(0,1000) for _ in range(10)]

print(a)
# [326, 859, 181, 700, 582, 144, 71, 495, 381, 838]

print(heapq.nlargest(2, a))
# [859, 838]

print(heapq.nsmallest(2, a))
# [71, 144]

ただし注意もあり,公式ドキュメントによれば,基本的に一回だけ大きい/小さい順に値を取り出すだけであれば,
配列をソートして取り出した方が高速です.
また最大値/最小値を取り出すのもmax / minの方が高速です.
この処理を複数回行う際に,iterable objectをヒープ化して実行するのが良い,とのことでした.

参考

heapq --- ヒープキューアルゴリズム — Python 3.10.0b2 ドキュメント