python TopKアルゴリズム

TopKアルゴリズム
配列の中の最小のk個数を探して、topk問題とも言います.
このアルゴリズムが解決すべき問題は,線形時間内に無秩序シーケンスのkk番目に大きい数を見つけることである.
例えば、n個の整数を入力し、その中で最小のK個の数を探し出す.例えば4,5,1,6,2,7,3,8の8つの数字を入力すると、最小の4つの数字は1,2,3,4である.
考え方:
高速ソートのpartition()メソッドは、a[l...j-1]がa[j]以下であり、a[j+1...h]がa[j]以上である整数jを返し、このときa[j]は配列のj番目の大きな要素である.この特性を用いて配列のK番目の要素を見つけることができ,このK番目の要素を探すアルゴリズムを高速選択アルゴリズムと呼ぶ.
 

# -*- coding: gbk -*-
def partition(seq):
    pi, seq = seq[0], seq[1:]                 #        
    lo = [x for x in seq if x <= pi]#             
    hi = [x for x in seq if x > pi]##             
    return lo, pi, hi

def select(seq, k):
    lo, pi, hi = partition(seq)
    m = len(lo)#            
    if m == k: return pi
    if m < k: return select(hi, k-m-1)
    return select(lo, k)

if __name__ == '__main__':
    seq=(1,2,3,4,5)
    print(partition(seq))
    print(select(seq,3))