python实现堆（最大堆、最小堆、最小最大堆）

原创

yuezht

修改于 2023-03-30 22:21:47

2K1

修改于 2023-03-30 22:21:47

1. 最大堆

class MaxHeap:
    def __init__(self):
        self.heap = []

    def parent(self, i):
        return (i - 1) // 2

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def get_max(self):
        if not self.heap:
            return None
        return self.heap[0]

    def insert(self, item):
        self.heap.append(item)
        self._heapify_up(len(self.heap) - 1)

    def extract_max(self):
        if not self.heap:
            return None

        max_item = self.heap[0]
        last_item = self.heap.pop()
        if self.heap:
            self.heap[0] = last_item
            self._heapify_down(0)
        return max_item

    def _heapify_up(self, i):
        while i > 0 and self.heap[i] > self.heap[self.parent(i)]:
            self.heap[i], self.heap[self.parent(i)] = self.heap[self.parent(i)], self.heap[i]
            i = self.parent(i)

    def _heapify_down(self, i):
        max_index = i
        left = self.left_child(i)
        if left < len(self.heap) and self.heap[left] > self.heap[max_index]:
            max_index = left
        right = self.right_child(i)
        if right < len(self.heap) and self.heap[right] > self.heap[max_index]:
            max_index = right
        if i != max_index:
            self.heap[i], self.heap[max_index] = self.heap[max_index], self.heap[i]
            self._heapify_down(max_index)


if __name__ == '__main__':

    max_heap = MaxHeap()
    max_heap.insert(1)
    max_heap.insert(2)
    max_heap.insert(0)
    max_heap.insert(8)
    print(max_heap.get_max())

2. 最小堆

class MinHeap:
    def __init__(self):
        self.heap = []

    def parent(self, i):
        return (i - 1) // 2

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def get_min(self):
        if not self.heap:
            return None
        return self.heap[0]

    def insert(self, item):
        self.heap.append(item)
        self._heapify_up(len(self.heap) - 1)

    def extract_min(self):
        if not self.heap:
            return None

        min_item = self.heap[0]
        last_item = self.heap.pop()
        if self.heap:
            self.heap[0] = last_item
            self._heapify_down(0)
        return min_item

    def _heapify_up(self, i):
        while i > 0 and self.heap[i] < self.heap[self.parent(i)]:
            self.heap[i], self.heap[self.parent(i)] = self.heap[self.parent(i)], self.heap[i]
            i = self.parent(i)

    def _heapify_down(self, i):
        min_index = i
        left = self.left_child(i)
        if left < len(self.heap) and self.heap[left] < self.heap[min_index]:
            min_index = left
        right = self.right_child(i)
        if right < len(self.heap) and self.heap[right] < self.heap[min_index]:
            min_index = right
        if i != min_index:
            self.heap[i], self.heap[min_index] = self.heap[min_index], self.heap[i]
            self._heapify_down(min_index)

3. 最小-最大堆

最小-最大堆的性质是：树中偶数层的每个节点都小于它的所有后代，而树中奇数层的每个节点都大于它的所有后代。

用途双端优先级队列

class MinMaxHeap:
    def __init__(self):
        self.heap = []

    def parent(self, i):
        return (i - 1) // 2

    def left_child(self, i):
        return 2 * i + 1

    def right_child(self, i):
        return 2 * i + 2

    def get_min(self):
        if not self.heap:
            return None
        return self.heap[0]

    def get_max(self):
        if not self.heap:
            return None
        if len(self.heap) == 1:
            return self.heap[0]
        if len(self.heap) == 2:
            return self.heap[1] if self.heap[1] > self.heap[0] else self.heap[0]
        return self.heap[1] if self.heap[1] > self.heap[2] else self.heap[2]

    def insert(self, item):
        self.heap.append(item)
        self._heapify_up(len(self.heap) - 1)

    def extract_min(self):
        if not self.heap:
            return None

        min_item = self.heap[0]
        last_item = self.heap.pop()
        if self.heap:
            self.heap[0] = last_item
            self._heapify_down_min(0)
        return min_item

    def extract_max(self):
        if not self.heap:
            return None

        max_item = self.get_max()
        max_index = self.heap.index(max_item)
        self.heap[max_index] = self.heap[-1]
        self.heap.pop()
        if max_index < len(self.heap):
            self._heapify_down_max(max_index)
        return max_item

    def _heapify_up(self, i):
        if i == 0:
            return
        parent = self.parent(i)
        if self.heap[i] < self.heap[parent]:
            self.heap[i], self.heap[parent] = self.heap[parent], self.heap[i]
            self._heapify_up_max(parent)
        else:
            self._heapify_up_min(i)

    def _heapify_up_min(self, i):
        grandparent = self.parent(self.parent(i))
        if i > 2 and self.heap[i] < self.heap[grandparent]:
            self.heap[i], self.heap[grandparent] = self.heap[grandparent], self.heap[i]
            self._heapify_up_min(grandparent)

    def _heapify_up_max(self, i):
        grandparent = self.parent(self.parent(i))
        if i > 2 and self.heap[i] > self.heap[grandparent]:
            self.heap[i], self.heap[grandparent] = self.heap[grandparent], self.heap[i]
            self._heapify_up_max(grandparent)

    def _heapify_down_min(self, i):
        while True:
            min_index = i
            left = self.left_child(i)
            if left < len(self.heap) and self.heap[left] < self.heap[min_index]:
                min_index = left
            right = self.right_child(i)
            if right < len(self.heap) and self.heap[right] < self.heap[min_index]:
                min_index = right
            if i != min_index:
                self.heap[i], self.heap[min_index] = self.heap[min_index], self.heap[i]
                i = min_index
            else:
                break

    def _heapify_down_max(self, i):
        while True:
            max_index = i
            left = self.left_child(i)
            if left < len(self.heap) and self.heap[left] > self.heap[max_index]:
                max_index = left
            right = self.right_child(i)
            if right < len(self.heap) and self.heap[right] > self.heap[max_index]:
                max_index = right
            if i != max_index:
                self.heap[i], self.heap[max_index] = self.heap[max_index], self.heap[i]
                i = max_index
            else:
                break

在这个实现中，MinMaxHeap类代表一个min-max堆，包含一个list堆，用于存放堆中的元素。 parent、left_child 和right_child 方法分别返回节点的父节点、左子节点和右子节点的索引。 get_min 方法返回堆中的最小元素，get_max 方法返回堆中的最大元素。 insert 方法将一个元素插入到堆中并维护堆属性。 extract_min 方法从堆中移除最小元素并保持堆属性。 extract_max 方法从堆中移除最大元素并保持堆属性。

_heapify_up、_heapify_up_min、_heapify_up_max、_heapify_down_min 和 _heapify_down_max 方法用于维护最小-最大堆属性。 _heapify_up 在向堆中插入元素后调用，以确保元素位于正确的位置。 _heapify_up_min 和 _heapify_up_max 由 _heapify_up 调用以维护最小-最大堆属性。 _heapify_down_min 和 _heapify_down_max 分别被 extract_min 和 extract_max 调用，以维护 min-max 堆属性。

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

python