数据结构算法之美的实现与拓展：算法9堆排序

前端

2024-01-23 15:32:19

堆排序，又称堆积排序，是一种利用堆资料结构（Binary Heap）来实现的排序算法，具有时间复杂度为 O(n log n) 的特性。它被广泛应用于各种场景，包括数据库、网络、图形处理等领域。

在本文中，我们将详细探讨堆排序算法的原理与实现，并进一步拓展其应用，了解如何解决Top K问题和中位数问题。

堆排序算法

堆排序的工作原理是基于堆资料结构的。堆是一种特殊的完全二叉树，其中每个节点的值都大于或小于其左右子节点的值。这种结构使得堆具有某些特殊的性质，使得我们可以利用它来实现排序。

堆排序算法分为以下几个步骤：

将待排序的元素构建成一个堆。
将堆顶元素与最后一个元素交换，并重新调整堆结构。
重复步骤 2，直到堆中只剩下一个元素。

下面我们详细解释一下每个步骤。

构建堆

构建堆的过程称为堆化（Heapify）。堆化算法从最后一个非叶节点开始，依次对每个非叶节点及其子节点进行比较，将较大的元素交换到父节点的位置，并依次交换下去，直到满足堆的性质。

def heapify(arr, n, i):
    largest = i
    left = 2 * i + 1
    right = 2 * i + 2

    if left < n and arr[left] > arr[largest]:
        largest = left

    if right < n and arr[right] > arr[largest]:
        largest = right

    if largest != i:
        arr[i], arr[largest] = arr[largest], arr[i]
        heapify(arr, n, largest)

排序

堆化完成之后，堆顶元素就是待排序元素中的最大值。我们将堆顶元素与最后一个元素交换，并重新调整堆结构，此时堆顶元素就是待排序元素中的次大值。重复这个过程，直到堆中只剩下一个元素，即完成排序。

def heap_sort(arr):
    n = len(arr)

    for i in range(n // 2 - 1, -1, -1):
        heapify(arr, n, i)

    for i in range(n - 1, 0, -1):
        arr[i], arr[0] = arr[0], arr[i]
        heapify(arr, i, 0)

算法拓展

除了排序之外，堆还可用来解决许多其他问题，例如Top K问题和中位数问题。

Top K问题

Top K问题是指在给定一个数组中，找到最大的或最小的 K 个元素。我们可以利用堆来解决Top K问题。

一种方法是构建一个小顶堆，将前 K 个元素插入堆中，然后依次插入剩余的元素。如果新插入的元素比堆顶元素大，则将堆顶元素弹出，并插入新元素。这样，堆中始终保存着最大的 K 个元素。

def top_k(arr, k):
    heap = []
    for i in range(k):
        heapq.heappush(heap, arr[i])

    for i in range(k, len(arr)):
        if arr[i] > heap[0]:
            heapq.heappop(heap)
            heapq.heappush(heap, arr[i])

    return heap

另一种方法是构建一个大顶堆，将前 K 个元素插入堆中，然后依次插入剩余的元素。如果新插入的元素比堆顶元素小，则将堆顶元素弹出，并插入新元素。这样，堆中始终保存着最小的 K 个元素。

def top_k_min(arr, k):
    heap = []
    for i in range(k):
        heapq.heappush(heap, -arr[i])

    for i in range(k, len(arr)):
        if -arr[i] > heap[0]:
            heapq.heappop(heap)
            heapq.heappush(heap, -arr[i])

    return [-x for x in heap]

中位数问题

中位数问题是指在给定一个数组中，找到一个元素，使得数组中有一半的元素大于或等于它，另一半的元素小于或等于它。我们可以利用堆来解决中位数问题。

一种方法是构建一个小顶堆和一个大顶堆。小顶堆中保存着数组中较大的元素，大顶堆中保存着数组中较小的元素。如果两个堆的大小相等，则中位数就是两个堆顶元素的平均值。否则，如果小顶堆的大小大于大顶堆的大小，则中位数就是小顶堆的堆顶元素。否则，中位数就是大顶堆的堆顶元素。

def median(arr):
    n = len(arr)
    if n % 2 == 0:
        small = []
        large = []
        for i in range(n):
            if i < n // 2:
                heapq.heappush(small, -arr[i])
            else:
                heapq.heappush(large, arr[i])

        while len(small) > len(large):
            heapq.heappush(large, -heapq.heappop(small))

        while len(large) > len(small) + 1:
            heapq.heappush(small, -heapq.heappop(large))

        return (heapq.heappop(small) - heapq.heappop(large)) / 2
    else:
        small = []
        for i in range(n):
            heapq.heappush(small, -arr[i])

        return -heapq.heappop(small)