Python三种方法实现topk问题(源码)

# topK问题  数组中有n个元素求前k个最大的数
# 1. 快排或小顶堆排n个数 返回前k个数  --- 时复为O(n+nlog_2n+k)
# 2. 第一次优化：首先根据n数组建立一个大顶堆 每次获取arr[0](并将其移除) 原地移除的方法是将arr[0]与arr[-1]对调 后在arr[0:-1)时向下调整法 反复上述步骤 直至k次 则获得了前k个最大的数  ------时复为O(n + klog_2n) 前一个n是建堆的时复，后面是进行了k次向下调整法，这样则当n很大时 log_2n趋于稳定 此时为线性复杂度O(n)
# 3. 第二次优化：当n很大时，上述时间复杂度还是很高， 还可以如何做？ 若是求前k个最大的数，则建元素个数为k的小堆，每加入一个元素与堆顶元素比较，若比堆顶大则将堆顶元素替换，再做向下调整法，直至遍历完数组，那么小堆中维护的就是数组中前k个大的数。时复最差为O(k+(n-k)log_2n)前面是k个堆的排序 后面是最糟糕的情况 即后面每次都要调整。

第一次优化代码

# 向下调整算法
def ajustDown(arr, end, parent_index):
    """
    :param arr: 需要调整的arr数组
    :param end: 需要调整的arr数组的末尾索引+1
    :param parent_index: 从该节点开始调整
    :return: 原地更改数组 无返回值
    """
    # 针对arr只有一个元素的情况 不需处理
    if parent_index < 0:
        return
    # 先认为当前父节点为最小
    smallest_index = parent_index
    # 获取左右孩子下标
    leftchild_index = 2 * parent_index + 1
    rightchild_index = 2 * parent_index + 2
    # 找出父 左 右 孩子最小的index
    if leftchild_index < end and arr[leftchild_index] < arr[parent_index]:
        smallest_index = leftchild_index
    if rightchild_index < end and arr[rightchild_index] < arr[smallest_index]:
        smallest_index = rightchild_index
    # 如果更改了下标 则进行交换 并递归更新
    if not (smallest_index == parent_index):
        arr[smallest_index], arr[parent_index] = arr[parent_index], arr[smallest_index]
        ajustDown(arr, end, smallest_index)

if __name__ == '__main__':
    # 测试用例
    # ---------- 对[整个数组 因此end=len(arr)]直接建堆 --------------
    arr = [2,11,8,5,4,9,35,36,1,6]
    for i in range(len(arr), -1, -1):
        parent_index = (i - 1) // 2
        ajustDown(arr, len(arr), parent_index)
    print(arr)
    # ------------ 对数组进行堆排序 ---------------
    # 堆排序分为两步 先对数组建堆 将头(最小元素)尾对调 对头进行向下调整法(这时不可以包括尾) 复原堆的位置
    # 这里建的小顶堆 因此是升序排序
    arr2 = [1, 2, 8, 5, 4, 9, 35, 36, 11, 6]
    # 建堆
    for i in range(len(arr2), -1, -1):
        parent_index = (i - 1) // 2
        ajustDown(arr2, len(arr2), parent_index)
    print(arr2)
    # 对调 及 对头 进行向下调整法
    for tail in range(len(arr2)-1, 0, -1):
        arr2[0], arr2[tail] = arr2[tail], arr2[0]
        ajustDown(arr2, tail, 0)
    print(arr2)

# 执行结果
# [1, 2, 8, 5, 4, 9, 35, 36, 11, 6]
# [1, 2, 8, 5, 4, 9, 35, 36, 11, 6]
# [36, 35, 11, 9, 8, 6, 5, 4, 2, 1]

第二次优化与第一次类似 只是换成大顶堆去实现

第三次优化如下：

# 向下调整算法
def ajustDown(arr, end, parent_index):
    """
    :param arr: 需要调整的arr数组
    :param end: 需要调整的arr数组的末尾索引+1
    :param parent_index: 从该节点开始调整
    :return: 原地更改数组 无返回值
    """
    # 针对arr只有一个元素的情况 不需处理
    if parent_index < 0:
        return
    # 先认为当前父节点为最小
    smallest_index = parent_index
    # 获取左右孩子下标
    leftchild_index = 2 * parent_index + 1
    rightchild_index = 2 * parent_index + 2
    # 找出父 左 右 孩子最小的index
    if leftchild_index < end and arr[leftchild_index] < arr[parent_index]:
        smallest_index = leftchild_index
    if rightchild_index < end and arr[rightchild_index] < arr[smallest_index]:
        smallest_index = rightchild_index
    # 如果更改了下标 则进行交换 并递归更新
    if not (smallest_index == parent_index):
        arr[smallest_index], arr[parent_index] = arr[parent_index], arr[smallest_index]
        ajustDown(arr, end, smallest_index)

if __name__ == '__main__':
    # 给定n个元素的arr个数组，求前k个最大的数，利用小顶堆--始终维护遍历到的最大的k个数，则将数组遍历完毕，就得到整个数组的前k个最大的数
    arr = [1,5,1,3,6,9,4,5,6,8,1,2,5,63,2,46,52,1,6,25,3,56,265,235,1,5,36,523]
    k = 5
    # 给数组arr前k个数进行建小顶堆
    for i in range(k-1, -1, -1):
        parent_index = (i - 1) // 2
        ajustDown(arr, k, parent_index)
    # 从数组的索引k进行遍历 遇到大于堆顶 则跟堆顶替换 并用向下调整法 始终维护遍历到当前位置的前k个最大的数
    for i in range(k,len(arr)):
        if arr[i] > arr[0]:
            arr[i], arr[0] = arr[0], arr[i]
            ajustDown(arr, k, 0)
    # 此时数组前k个数 为最大的数 输出结果
    print(arr[:k])
# 执行结果 [56, 63, 523, 265, 235]