js 实现堆

156 阅读1分钟

通过传入type为 ‘small’ 或者‘big’配置小顶堆或大顶堆,通过cmp自定义比较函数可以使堆适用多种场景

class Heap {
   constructor(type,cmp) {
       this.type = type || 'big';
       this.h = [];
       this.cmp = cmp ? cmp : (a, b)=>{
           if(this.type === 'big'){
               return a > b;
           }else {
               return a < b;
           }
       }
   }

   swap(i, k) {
       [this.h[i],this.h[k]] = [this.h[k], this.h[i]];
   }

   push(val) {
       this.h.push(val);

       // 向上调整
       let curIndex = this.h.length - 1;
       let parentIndex = Math.floor((curIndex-1) / 2 );
       while (parentIndex >= 0){
           if(this.cmp(this.h[parentIndex],this.h[curIndex])) break;
           this.swap(curIndex,parentIndex);
           curIndex = parentIndex;
           parentIndex = Math.floor((curIndex-1) / 2 );
       }

   }

   pop() {
       if( this.empty() ) return null;
       this.swap(0, this.h.length - 1);
       const res = this.h.pop();
       // 向下调整
        let curIndex = 0;
        let leftChildIndex = curIndex * 2 + 1;
        let temp = curIndex;
        while ( leftChildIndex <= this.size() -1 ) {
            if(this.cmp(this.h[leftChildIndex], this.h[curIndex])) temp = leftChildIndex;
            if((leftChildIndex +1 <= this.size() -1) 
            && this.cmp(this.h[leftChildIndex + 1], this.h[leftChildIndex]) 
            && this.cmp(this.h[leftChildIndex + 1], this.h[curIndex]) ) temp = leftChildIndex + 1;
            if(curIndex === temp) break;
            this.swap(curIndex, temp);
            curIndex = temp;
            leftChildIndex = curIndex * 2 + 1;
        }

       return res;
   }

   top() {
       return this.h[0];
   }

   empty() {
      return  this.h.length === 0;
   }

   size() {
       return this.h.length;
   }
}