[python] n个数中K个最小值

code:

# 类似于快排的思想，不同的地方在于每趟只需要往一个方向走
# 按照从小到大的顺序，寻找前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

最大堆
假设数组长度为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