Python中的最小堆实现

Question

这是我对在 Python 中使用 sift-down 实现最小堆的看法。 有什么办法可以缩短筛选代码？我尝试不使用那么多 if 子句，但我似乎无法让它更短。

import math


class MinHeap:
   def __init__(self, arr):
       self.arr = arr
       self.n = len(arr)

    def heapify(self):
        depth = int(math.log2(self.n))
        for i in range((2 ** depth) - 2, -1, -1):
            self.sift_down(i)

    def sift_down(self, i):
        while (2 * i) + 1 < self.n or (2 * i) + 2 < self.n:
            i = self.sift_down_level(i)

    def sift_down_level(self, i):
        left_smaller = right_smaller = False
        if self.is_left_child_exists(i) and self.arr[(2 * i) + 1] < self.arr[i]:
            left_smaller = True
        if self.is_right_child_exists(i) and self.arr[(2 * i) + 2] < self.arr[i]:
            right_smaller = True

        if left_smaller and right_smaller:
            if self.arr[(2 * i) + 1] < self.arr[(2 * i) + 2]:
                self.arr[(2 * i) + 1], self.arr[i] = self.arr[i], self.arr[(2 * i) + 1]
                return (2 * i) + 1
            else:
                self.arr[(2 * i) + 2], self.arr[i] = self.arr[i], self.arr[(2 * i) + 2]
                return (2 * i) + 2
        elif left_smaller:
            self.arr[(2 * i) + 1], self.arr[i] = self.arr[i], self.arr[(2 * i) + 1]
            return (2 * i) + 1
        elif right_smaller:
            self.arr[(2 * i) + 2], self.arr[i] = self.arr[i], self.arr[(2 * i) + 2]
            return (2 * i) + 2
        return -1

    def is_left_child_exists(self, i):
        if (2 * i) + 1 < self.n:
            return True
        return False

    def is_right_child_exists(self, i):
        if (2 * i) + 2 < self.n:
            return True
        return False


nums = [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
heap = MinHeap(nums)
heap.heapify()

Answer 1

您应该避免代码重复：只在一个地方执行交换。避免多次重新计算2 * i + 1。而且，是的，减少您检查的条件数量。您只需要检查两个孩子中的哪一个（如果有两个）的值最小，然后将 that 的值与父级的值进行比较。就是这样。

其他一些评论：

• 我不会将数组的长度存储在属性中，因为这样每次从堆中推入或拉出值时，您都需要保持更新的开销。
• 无需执行对数和幂运算。只需计算最后一个节点的父节点在哪里。最后一个节点位于索引len()-1，因此其父节点位于(len()-1-1)//2，即len()//2-1。
• sift_down 中的while 条件不应重复您将在sift_down_level 中执行的检查。只需验证返回值不是 -1。
• 并在初始化期间调用heapify

无需对您的设计进行其他更改，我们就能得到：

class MinHeap:
    def __init__(self, arr):
       self.arr = arr
       self.heapify()

    def heapify(self):
        for i in range(len(self.arr) // 2 - 1, -1, -1):
            self.sift_down(i)

    def sift_down(self, i):
        while i != -1:
            i = self.sift_down_level(i)

    def sift_down_level(self, i):
        child = i * 2 + 1
        if child < len(self.arr):
            if child + 1 < len(self.arr) and self.arr[child + 1] < self.arr[child]:
                child += 1
            if self.arr[child] < self.arr[i]:
                self.arr[child], self.arr[i] = self.arr[i], self.arr[child]
                return child
        return -1