带有伪代码的堆和优先级队列

Question

我在学校有一个问题，我需要在最小堆中找到第 k 个最小元素。所需的运行时间是 o(k^2)，我知道该怎么做。但如果我能把它降到 o(k*logk) 我得到奖金。我想过从最小堆做一个优先级队列，然后将堆的节点插入队列，然后将其取出，然后对最小堆的根的孩子做同样的事情，等等 k 次。我知道插入和弹出操作的时间复杂度是 o(logk)，因为优先级队列的初始大小是 1，并且在每 k 步中最多增加 1。因此，优先级队列中最多有 k+1 个元素。

我了解我需要做什么，但发现用伪代码实现它很复杂，任何想法或指南都会很棒。

谢谢

Answer 1

是的，你的想法听起来不错。我将展示如何在 Python 中执行此操作。

从数组构建堆

所以，首先在输入数组的顶部构建一个最小堆：

def create_min_heap(array):
    min_heap = []
    for value in array:
        heappush(min_heap, value)

    return min_heap

从堆中计算 k min 个元素

创建一个辅助最小堆，用于在 O(klogk) 中检索所有 k 最小元素。在每一步，都会从中弹出一个最小元素，并添加其2 子节点（子节点可以在原始最小堆中找到）。请注意，将其他节点的子节点添加到辅助最小堆是没有意义的，因为它们不能小于其父节点（每个堆属性）。

def k_min_elements(min_heap, k):
    result = list()
    helper_min_heap = []
    heappush(helper_min_heap, (min_heap[0],0))
    while len(result) < k:
        min_node = heappop(helper_min_heap)
        value = min_node[0]
        index = min_node[1]

        left_index = index*2 + 1
        right_index = left_index + 1
        if left_index < len(min_heap):
            heappush(helper_min_heap, (min_heap[left_index], left_index))
        if right_index < len(min_heap):
            heappush(helper_min_heap, (min_heap[right_index], right_index))

        result.append(value)

    return result 

完整代码

现在，完整的代码和示例输出。

from heapq import heappop
from heapq import heappush

def create_min_heap(array):
    min_heap = []
    for value in array:
        heappush(min_heap, value)

    return min_heap

def k_min_elements(min_heap, k):
    if k > len(min_heap) or k < 0:
        raise Exception("k is invalid")

    result = list()
    helper_min_heap = []
    heappush(helper_min_heap, (min_heap[0],0))
    while len(result) < k:
        min_node = heappop(helper_min_heap)
        value = min_node[0]
        index = min_node[1]

        left_index = index*2 + 1
        right_index = left_index + 1
        if left_index < len(min_heap):
            heappush(helper_min_heap, (min_heap[left_index], left_index))
        if right_index < len(min_heap):
            heappush(helper_min_heap, (min_heap[right_index], right_index))

        result.append(value)

    return result


min_heap = create_min_heap([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
print (k_min_elements(min_heap, 3))

[0, 1, 2]