python TopK算法

TopK算法

寻找数组中的最小的k个数，也叫topk问题。

该算法要解决的问题是：在线性时间内找到一个无序序列中第 kk 大的数。

如：输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

思路：

快速排序的 partition() 方法，会返回一个整数 j 使得 a[l…j-1] 小于等于 a[j]，且 a[j+1…h] 大于等于 a[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))