一文学会数据结构-堆

什么是堆

在计算机科学中，堆是一个特殊的树状数据结构。堆通常有如下几个特性：

1. 堆是完全二叉树；
2. 堆中每个节点的值都大于或等于（小于或等于）其子节点的值，这个取值规则称为堆的“属性”;
3. 堆顶元素（即根节点）总是为最大值或最小值。

堆的种类

堆分为小根堆和大根堆两种。小根堆要求每个节点的值都不大于其父节点的值，即A[PARENT[i]] <= A[i]；大根堆要求每个节点的值都不小于其父节点的值，即A[PARENT[i]] >= A[i]。

优先队列

堆可以用于组成优先队列。优先队列是一种数据结构，其中每个元素有一个“优先级”或“关键值”，优先级高的元素先被处理。当然，堆并不是唯一可以组成优先队列的数据结构，但堆的效率要高得多。

实现堆

堆是通过数组实现的。令i表示二叉树中节点的序号，A[i]表示存储在该节点的元素。

1. 父节点的位置为 Parent(i) = floor(i/2)
2. 左子节点的位置为 Left(i) = 2i
3. 右子节点的位置为 Right(i) = 2i + 1

实现中我们可以使用数组存储堆，这样可以快速寻找到每个节点的父节点和子节点。

堆的操作

下面介绍堆的几个常见操作：

堆化（Heapify）

将i结点以下的堆整体调整成最大堆或最小堆，时间复杂度为O(logN)。

# 以大根堆为例
def heapify(arr, n, i):
    # 初始化最大值位置
    largest = i 
    left = 2 * i + 1
    right = 2 * i + 2

    # 如果左子节点的值大于父节点的值，更新最大值位置
    if left < n and arr[i] < arr[left]:
        largest = left

    # 如果右子节点的值大于最大值位置节点的值，更新最大值位置
    if right < n and arr[largest] < arr[right]:
        largest = right

    # 如果最大值位置不是父节点位置，则交换这两个节点的值，并递归调用heapify
    if largest != i:
        arr[i], arr[largest] = arr[largest], arr[i]
        heapify(arr, n, largest)

建堆

通过对所有元素进行一次heapify操作实现，建堆的时间复杂度为O(n)。

def build_heap(arr, n):
    # 从最后一个非叶子节点开始heapify
    for i in range(int(n / 2) - 1, -1, -1):
        heapify(arr, n, i)

插入（Insertion）

将新元素插入到堆中，保证堆的属性。时间复杂度为O(logN)。

def insert(arr, element):
    arr.append(element)
    i = len(arr) - 1
    p = int((i - 1) / 2)
    while i > 0 and arr[p] < arr[i]:
        arr[p], arr[i] = arr[i], arr[p]
        i = p
        p = int((i - 1) / 2)

删除堆顶元素（Remove）

删除堆顶元素并返回其值，时间复杂度为O(logN)。

def remove(arr):
    n = len(arr)
    # 如果堆为空，则返回None
    if n == 0:
        return None
    # 如果堆中只有一个元素，则直接返回该元素
    if n == 1:
        return arr.pop()
    # 弹出堆顶元素，并将末尾元素移动到堆顶
    top = arr[0]
    arr[0] = arr.pop()
    n = len(arr)
    # 对新的堆顶元素进行一次heapify操作，使得剩余元素保持堆的属性
    heapify(arr, n, 0)
    return top

示例

示例1

给定数组 [3, 2, 1, 6, 5, 4]，按升序排列后的结果为 [1, 2, 3, 4, 5, 6]。下面是一个Python实现：

# 堆排序
def heap_sort(arr):
    # 首先将数组构建为最大堆
    build_heap(arr, len(arr))
    # 将最大的元素交换到数组末尾，并缩小堆的范围
    for i in range(len(arr) - 1, 0, -1):
        arr[i], arr[0] = arr[0], arr[i]
        heapify(arr, i, 0)

arr = [3, 2, 1, 6, 5, 4]
heap_sort(arr)
print(arr)

输出： [1, 2, 3, 4, 5, 6]。

示例2

使用堆来解决“TOP K 问题”，即在一个无序的整数数组中，找到最大的K个数。下面是一个Python实现：

import heapq

def top_k(arr, k):
    # 取前K个数，建立一个小根堆
    heap = arr[:k]
    heapq.heapify(heap)
    # 对于剩下的元素，如果比堆顶元素大，则弹出堆顶，并将该元素压入堆中
    for element in arr[k:]:
        if element > heap[0]:
            heapq.heappop(heap)
            heapq.heappush(heap, element)
    # 返回堆中的元素，即为前K大的数
    return heap

arr = [4, 3, 6, 5, 9, 8]
print(top_k(arr, 3)) # 输出: [6, 8, 9]

输出 [6, 8, 9]，为原数组中最大的三个数。