class Heap: def __init__(self, is_min_heap=True): self.heap = [] self.is_min_heap = is_min_heap def _compare(self, a, b): """比较函数，根据是最小堆还是最大堆决定比较方式""" if self.is_min_heap: return a < b else: return a > b def _parent(self, index): """获取父节点索引""" return (index - 1) // 2 def _left_child(self, index): """获取左子节点索引""" return 2 * index + 1 def _right_child(self, index): """获取右子节点索引""" return 2 * index + 2 def _swap(self, i, j): """交换堆中两个元素""" self.heap[i], self.heap[j] = self.heap[j], self.heap[i] def _heapify_up(self, index): """向上调整堆结构""" parent = self._parent(index) if index > 0 and self._compare(self.heap[index], self.heap[parent]): self._swap(index, parent) self._heapify_up(parent) def _heapify_down(self, index): """向下调整堆结构""" largest = index left = self._left_child(index) right = self._right_child(index) if left < len(self.heap) and self._compare(self.heap[left], self.heap[largest]): largest = left if right < len(self.heap) and self._compare(self.heap[right], self.heap[largest]): largest = right if largest != index: self._swap(index, largest) self._heapify_down(largest) def insert(self, value): """插入元素到堆中""" self.heap.append(value) self._heapify_up(len(self.heap) - 1) def extract_top(self): """提取堆顶元素""" if len(self.heap) == 0: return None top = self.heap[0] self._swap(0, len(self.heap) - 1) self.heap.pop() self._heapify_down(0) return top def get_top(self): """获取堆顶元素""" if len(self.heap) == 0: return None return self.heap[0] def size(self): """获取堆的大小""" return len(self.heap) def is_empty(self): """判断堆是否为空""" return len(self.heap) == 0 # 测试代码 if __name__ == "__main__": # 测试最小堆 min_heap = Heap(is_min_heap=True) print("=== 测试最小堆 ===") data = [5, 3, 8, 1, 2, 7] for num in data: min_heap.insert(num) print(f"堆顶元素: {min_heap.get_top()}") print(f"堆大小: {min_heap.size()}") print("提取堆顶元素:") while not min_heap.is_empty(): print(min_heap.extract_top(), end=" ") print() # 测试最大堆 max_heap = Heap(is_min_heap=False) print("=== 测试最大堆 ===") data = [5, 3, 8, 1, 2, 7] for num in data: max_heap.insert(num) print(f"堆顶元素: {max_heap.get_top()}") print(f"堆大小: {max_heap.size()}") print("提取堆顶元素:") while not max_heap.is_empty(): print(max_heap.extract_top(), end=" ") print()
