计算机算法-线性寻找第k小（大）元素（Top K）BFPRT算法

问题描述

从n个元素的序列中选出第k大（第k小）的元素（求第K大就是求第n-K+1小），BFPRT可以保证在最坏情况下仍为线性时间复杂度。整体的思想与快速排序思想相似，借助了分治的思想，以及快排的分区间思想。

算法步骤

（1）对数组元素进行分组，每组5个元素，剩余不足五个元素的为一组；
（2）找出每一组的中位数，这里可以用任意的排序方法，因为只有五个元素，所以复杂度为常数级；
（3）递归调用上一个算法找到中位数中的中位数，一直找到只剩一个元素，记为m；
（4）用m来分割数组，使m前面的都比m小，后面的都比m大。
（5）判断m下标和k的大小，来确定目标元素位置；
（6）确定位置之后再次进行递归；

代码实现

def merge(lst, l, mid, r):
    left = lst[l:mid + 1]
    right = lst[mid + 1:r + 1]  #
    k = l
    while left and right:
        if left[0] <= right[0]:
            lst[k] = left.pop(0)
        else:
            lst[k] = right.pop(0)
        k += 1
    tail = left if left else right
    for n in tail:
        lst[k] = n
        k += 1


def mergesort(lst, l, r):
    if l < r:
        mid = (l + r - 1) // 2
        mergesort(lst, l, mid)
        mergesort(lst, mid + 1, r)
        merge(lst, l, mid, r)
    return lst


def partion(a, m, m_index):
    # 对a进行排序，使得比m小的元素放在m前面，比m大的元素放在m后面
    # 输入：m_index（m在a中的index）
    # 返回m前面, m后面各自元素的数目,以及m在新数组中的index

    # 将m与数组第一个元素交换位置，然后即可像快速排序一样将所有元素以m为中间元素
    # 分到左右两边
    tmp = a[0]
    a[0] = m
    a[m_index] = tmp

    i = 0
    j = len(a) - 1
    control_m = a[0]
    while i < j:
        while i < j and a[j] >= control_m:
            j -= 1

        a[i] = a[j]

        while i < j and a[i] <= control_m:
            i += 1
        a[j] = a[i]

    # 此时i = j， a[i]应该是最终的控制关键字所在位置
    a[i] = control_m
    print("m:{},after partion, a:{}".format(m, a))

    return i, len(a) - i - 1, i


def Bfprt(a, k):
    # 得到a中第k大的元素

    if len(a) < 5:
        # 元素数目不足5个时，直接排序，之后取下表为k-1的数,即为第k大的元素
        mergesort(a, 0, len(a) - 1)
        return a[k - 1]

    total_num = len(a)
    splits = total_num // 5  # 一共分成这么多组  //出来是整数  /为float型
    #print(splits)
    # 获取每一个分组的中位数
    split_medians = []
    for i in range(splits):
        cur = mergesort(a[i * 5:(i + 1) * 5], 0, 4)
        #print (cur)
        mid = cur[2]
        split_medians.append(mid)

    # 递归调用Bfprt算法，求这些中位数的中第splits//2大的元素，也就是中位数的中位数
    m = Bfprt(split_medians, splits // 2)
    # 求出m在a中的index
    m_index = [i for i in range(total_num) if a[i] == m][0]

    # 根据m对a进行划分，线性复杂度
    # num1， num2:小于m的元素的数目， 大于m的元素数目
    num1, num2, m_index = partion(a, m, m_index)
    if k == num1 + 1:#第k大为m
        return m
    elif k <= num1:
        # 在s1集合中
        return Bfprt(a[:m_index], k)
    else:
        # 在s2集合中
        return Bfprt(a[m_index + 1:], k - 1 - m_index)


a = [4, 2, 13, 8, 3, 10, 15, 18, 11, 1, 16, 17, 12, 19, 0, 6, 5, 7, 9, 14]
print("请输入k值：")
t = input()
k=int(t)
print(Bfprt(a, k))