找出数组中前K个最小的数-Python实现

最新推荐文章于 2024-04-19 17:08:46 发布

原创最新推荐文章于 2024-04-19 17:08:46 发布 · 8.2k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#Python #数据结构与算法

学习心得同时被 2 个专栏收录

3 篇文章

订阅专栏

学习札记

3 篇文章

订阅专栏

本文介绍两种高效算法来解决寻找数组中第K大元素的问题。一是通过构建二叉堆，实现元素筛选，时间复杂度为O(NLogK)。二是采用快排思想，直接定位目标元素，平均时间复杂度为O(N)。文章提供了详细的算法实现代码及分析。

  寻找数组中给定的第K大的数，或者前K个最大的数，与之同理，稍加改动即可

思路1：二叉堆。

假设数组长度为N，首先取前K个数，构建二叉堆（大顶堆），然后将剩余N-K个元素，依次与堆顶元素进行比较，若小于堆顶元素，则替换，并重新为大顶堆。

代码如下

# 最大堆下沉调整，始终保持最大堆
def downAdjust(ary_list, parent_index, length):
    tmp = ary_list[parent_index]
    child_index = 2 * parent_index + 1

    while child_index < length:
        if child_index + 1 < length and ary_list[child_index + 1] > ary_list[child_index]:
            child_index += 1

        if tmp >= ary_list[child_index]:
            break

        ary_list[parent_index] = ary_list[child_index]
        parent_index = child_index
        child_index = 2 * parent_index + 1

    ary_list[parent_index] = tmp
    pass

# 构建堆
def build_heap(ary_list, k):
    index = k // 2 - 1  # 最后一个非叶子结点
    while index >= 0:
        downAdjust(ary_list, index, k)
        index -= 1
    pass


# 利用最大堆找出前K个最小值
# 每次从原数组中拿出一个元素和当前堆顶值比较，
# 然后判断是否可以放入，放入后继续调整堆结构
def heapK(ary, nums, k):
    if nums <= k:
        return nums
        
    ks = ary[:k]
    build_heap(ks, k)           # 构建大顶堆（先不排序）
    # print('build heap:', ks)
    
    for index in range(k, nums):
        ele = ary[index]
        if ks[0] > ele:
            ks[0] = ele
            downAdjust(ks, 0, k)
            # print('heap adjust:', ks)

    # 如果需要则输出排序结果
    # heap_sort(ks)
    return ks
    pass


if __name__ == '__main__':

    # *** 测试方法1
    ary_list = [10, 2, 38, 9, 22, 53, 47, 7, 3, 97]
    nums = len(ary_list)
    print('{} original data:'.format(nums), ary_list)

    # # 原始数组的排列顺序（作为ks的对比）
    # build_heap(ary_list, nums)
    # heap_sort(ary_list)
    # print('{} original sorted data:'.format(nums), ary_list)

    for k in range(6, nums + 1):
        ks = heapK(ary_list, nums, k)
        print('{}th data:'.format(k), ks)
        break
    pass

结果如下图所示：

在这里插入图片描述

如果想要输出结果有序，则可以增加如下方法，进行堆排序即可。

# 堆排序（最大堆）
def heap_sort(ary):
    length = len(ary)

    index = length - 1
    # 依次移除堆顶元素（放入末尾），并将末尾元素放在堆顶，进行下沉调整，
    # 使得每次都会有非最大值上浮到堆顶，并重新调整为大顶堆；
    # 然后再重复上述操作。
    while index >= 0:
        tmp = ary[0]
        ary[0] = ary[index]
        ary[index] = tmp
        downAdjust(ary, 0, index)
        index -= 1

    pass

总结

针对二叉堆的思路，其实主要是寻找能容纳K个元素的容器，然后在该容器中进行筛选操作。
若想要输出有序结果，则可以选择不同的排序算法对K个元素进行排序即可。

时间复杂度

构建大顶堆 ：平均复杂度为O(KLogK)，；
元素筛选 ：剩余N-K个元素，最坏情况下每个元素都要进行堆调整，复杂度为O((N-K)LogK)；

所以总的平均时间复杂度为O(KLogK)+O((N-K)LogK)=O(NLogK)。

空间复杂度为O(K)

思路2：快排思想。

利用快排的思想，循环找到第K个位置安放正确的元素，此时K的左边是小于K位置元素的元素，右边是大于K位置元素的元素，即前K个元素就是问题答案。

代码如下：

# 类似于快排的思想，不同的地方在于每趟只需要往一个方向走
# 按照从小到大的顺序，寻找前K个最小值
def qselect(ary_list, k):
    if len(ary_list) < k:
        return ary_list

    tmp = ary_list[0]
    left = [x for x in ary_list[1:] if x <= tmp] + [tmp]
    llen = len(left)
    if llen == k:
        return left
    if llen > k:
        return qselect(left, k)
    else:
        right = [x for x in ary_list[1:] if x > tmp]
        return left + qselect(right, k-llen)
    pass