数据结构：堆 heap

st_heap.py

# Created   : 2020-05-22 00:14:21
# 创建和调整：https://zhuanlan.zhihu.com/p/77527032
# 插入和删除：https://www.jianshu.com/p/e0a40d6748b8

import operator

default_op = operator.ge


def heap_adjust(A, i, size, op):
    """
    从当前节点向子节点延伸，一直到堆的限定大小，也就是不能修改堆大小范围外的数据
    """
    left = 2 * i + 1
    right = 2 * i + 2
    parent_index = i
    if left < size and op(A[left], A[parent_index]):
        parent_index = left
    if right < size and op(A[right], A[parent_index]):
        parent_index = right
    if parent_index != i:
        # print('.', end='')  # 输出的点数即置换次数
        A[parent_index], A[i] = A[i], A[parent_index]
        heap_adjust(A, parent_index, size, op=op)
    return A


def build_heap(A, size, op=default_op):
    for i in range(size // 2, -1, -1):
        heap_adjust(A, i, size, op)
    return A


def heap_sort(A, reverse=False):
    """默认升序"""
    size = len(A)
    op = operator.le if reverse else operator.ge
    A = build_heap(A, size, op=op)

    # 再把排好序的元素依次交换到目标位置
    for i in range(len(A) - 1, 0, -1):
        A[0], A[i] = A[i], A[0]
        A = heap_adjust(A, 0, i, op)
    return A


def heap_insert(A, x, op=default_op):
    A.append(x)
    current_index = len(A) - 1
    parent_index = current_index // 2
    while parent_index >= 0 and op(A[current_index], A[parent_index]):
        A[current_index], A[parent_index] = A[parent_index], A[current_index]
        if parent_index == 0:
            break
        current_index = parent_index
        parent_index = current_index // 2
    return A


def heap_pop(A, op=default_op):
    if not A:
        return None

    if len(A) == 1:
        return A.pop()

    rv = A[0]
    A[0] = A.pop()
    heap_adjust(A, 0, len(A), op)
    return rv


def top_k(nums, k):
    # 有限大小的最小堆
    op = operator.le
    b = build_heap(nums[:k], k, op)

    # 从k+1开始，依次比较，然后替换
    for i in range(k, len(nums)):
        if nums[i] > b[0]:
            nums[i], b[0] = b[0], nums[i]
            b = heap_adjust(b, 0, k, op)

    # 取最小值即第k大的值
    return b[0]


if __name__ == '__main__':
    print(heap_sort(list(range(99, 0, -1))))
    print(heap_sort(list(range(1, 100)), reverse=True))

    print(top_k(list(range(1, 10)), 3))
    print(top_k(list(range(1, 1000)), 4))

    b = build_heap(list(range(1, 10)), 9)
    print(b[0])
    heap_insert(b, 10)
    print(b[0])
    heap_insert(b, 11)
    print(b[0])

    for _ in range(len(b)+2):
        print(heap_pop(b))