[算法解腻]-大小顶堆

作者: 云峰 github: github.com/ihtml5

一、介绍

堆是一种特殊的树形结构。根据根节点的值与子节点的值的大小关系，堆又分为最大堆和最小堆。在最大堆中，每个节点的值总是大雨或等于其任意子节点的值，因此最大堆的根节点就是整个堆的最大值。在最小堆中，每个节点的值总是小于或等于任意子节点的值。因此最小堆的根节点是堆中的最小值。

二、应用场景

• 最大堆：常用于求解最小的k个数
• 最小堆：常用于求解频率最高的k个数

三、代码解腻

class Heap {
   constructor(data = [], compator) {
     this.data = data; // 传入初始顺序
     this.compator = compator; // 自定义对比函数，可实现最大堆和最小堆
     this.heapify(); // 堆排序
   }
   
   /*
     建立堆
   */
   heapify() {
     if (this.size() < 2) {
       return;
     }
     
     for (let i = 1; i < this.size() - 1; i++) {
        this.bubbleUp(i);
     }
   }
   /*
     获取堆大小
   */
   size() {
     return this.data.length;
   }
   /*
     追加堆
   */
   offer(value) {
     this.data.push(value);
     this.bubbleUp(this.size() - 1);
   }
   /*
     从堆顶移除元素
   */
   poll() {
     const last = this.data.pop(); // 弹出堆最后一个元素
     const result = this.data[0]; //  获取堆顶元素 
     
     if (this.size() !== 0 ) { // 堆不为空
        this.data[0] = last; // 将堆最后一个元素赋值给堆顶
        this.bubbleDown(0); // 进行向下堆排序
     }
     
     return result; // 返回堆顶元素
   }
   
   peek() {
     if (this.size() === 0) {
       return null; // 如果堆为空，返回null
     }
     return this.data[0];
     // 返回堆顶元素，和peek方法的差别是这里只获取元素的值，并不删除元素
   }
   
   /*
     向上堆排序
   */
   bubbleUp(index) {
     while (index > 0) {
        const parentIndex = Math.floor((index - 1) / 2); // 获取父元素索引
        if (this.compator(this.data[index], this.data[parentIndex])) {
           this.swap(index, parentIndex); // 交换元素
           index = parentIndex; // 更新索引值，继续向上排序
        } else {
           break;
        }
     }
   }
   /*
     向下堆排序
   */
   bubbleDown(index) {
     const lastIndex = this.size() - 1; // 堆最后一个元素索引
     while (true) {
       const leftIndex = 2 * index + 1; // 左子元素索引
       const rightIndex = 2 * index + 2; // 右子元素索引
       let findIndex = index;
       if (leftIndex <= lastIndex && this.compator(this.data[leftIndex], this.data[findIndex])) {
         findIndex = leftIndex; // 如果左元素小于或者大于findIndex，交换元素
       }
       
       if (rightIndex <= lastIndex && this.compator(this.data[rightIndex], this.data[findIndex])) {
         findIndex = rightIndex; // 如果右元素小于或者大于findIndex，交换元素
       }
       if (index !== findIndex) {
         this.swap(index, findIndex); // 交互index和findIndex元素
         index = findIndex;
       } else {
         break; // 终止遍历
       }
     }
   }
    /*
     交换元素
   */
   swap(index1, index2) {
     [this.data[index1], this.data[index2]] = [this.data[index2], this.data[index1]];
   }
}

四、测试用例

// 最小堆
const minheap = new Heap([], (a, b) => a - b < 0);
minheap.offer(2);
minheap.offer(3);
minheap.offer(-1);
console.log(minheap.poll());
minheap.offer(-3);
console.log(minheap.peek());

// 最大堆
const minheap = new Heap([], (a, b) => a - b > 0);
minheap.offer(2);
minheap.offer(3);
minheap.offer(-1);
console.log(minheap.poll());
minheap.offer(-3);
console.log(minheap.peek());