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

  寻找数组中给定的第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个元素进行排序即可。

时间复杂度

1. 构建大顶堆 ：平均复杂度为O(KLogK)，；
2. 元素筛选 ：剩余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

结果如下

时间复杂度：最坏情况下，每次只能找出一个最小值，总共需要找K次，复杂度为O(KN)；最好情况下只需要遍历依次即可找到，复杂度为O(N)；当K远小于N时，平均时间复杂度为O(N)。

空间复杂度：因为开辟了存储待寻找元素的数组，所以空间复杂度为O(N)。

总结
关于此类型问题，目前所总结和学习到的以这两种方案为主，仅作分享。
若有错漏，欢迎交流指正。

参考链接:
（1）漫画：什么是堆排序？.
（2）最大K个数问题的Python版解法总结

感谢~