JavaScript로 힙(Heap) 자료구조 구현하기

힙(Heap) 이란?

힙 요소 추가

1. 트리의 마지막 노드에 요소를 추가한다.

2. 추가된 노드가 부모 노드보다 우선순위가 높다면(낮다면) 부모 노드와 순서를 바꾼다.

3. 루트 노드까지 우선순위를 비교하여 2번 과정을 반복한다.

복사
push(item) {
  // 트리의 마지막 노드에 요소를 추가한다.
  this.heap.push(item);
  let currentNode = this.heap.length - 1;
  let parentNode = Math.floor(currentNode / 2);

  // 현재 노드와 부모 노드의 우선순위를 비교한다
  while (parentNode !== 0 && this.compare(this.heap[currentNode], this.heap[parentNode])) {
    // 우선순위가 높다면(낮다면) 부모 노드와 순서를 바꾼다.
    this._swap(parentNode, currentNode);
    // 현재 노드와 부모 노드의 인덱스를 업데이트한다.
    currentNode = parentNode;
    parentNode = Math.floor(currentNode / 2);
  }
}

힙 요소 제거

1. 루트 노드를 삭제한다.

2. 마지막 노드를 루트 노드에 위치시킨다.

3. 루트 노드와 자식 노드들 중에서 더 우선순위가 높은(낮은) 노드와 순서를 바꾼다.

4. 자식 노드들의 우선순위가 더 낮을(높을) 때 까지 3번 과정을 반복한다.

복사
pop() {
  if (this.heap.length === 1) return;
  if (this.heap.length === 2) return this.heap.pop();

  const top = this.heap[1];
  // 루트 노드를 삭제하고 마지막 노드를 루트 노드에 위치시킨다.
  this.heap[1] = this.heap.pop();

  let [currentNode, leftChild, rightChild] = [1, 2, 3];

  // 현재 노드와 자식 노드들의 우선 순위를 비교한다.
  while (
    this.compare(this.heap[leftChild], this.heap[currentNode]) ||
    this.compare(this.heap[rightChild], this.heap[currentNode])
  ) {
    // 자식 노드들 중에서 더 우선순위가 높은(낮은) 노드와 순서를 바꾼다.
    const targetNode = this.compare(this.heap[rightChild], this.heap[leftChild]) ? rightChild : leftChild;
    this._swap(currentNode, targetNode);
    // 현재 노드와 자식 노드들의 인덱스를 업데이트한다.
    currentNode = targetNode;
    leftChild = currentNode * 2;
    rightChild = currentNode * 2 + 1;
  }

  return top;
}

전체 코드

복사
class PriorityQueue {
  constructor(compare) {
    // 계산 편의성을 위해 0번 index는 사용하지 않는다.
    this.heap = [null];
    // 인스턴스 생성 시 우선순위를 비교하는 함수를 받는다.
    this.compare = compare;
  }

  push(item) {
    // 트리의 마지막 노드에 요소를 추가한다.
    this.heap.push(item);
    let currentNode = this.heap.length - 1;
    let parentNode = Math.floor(currentNode / 2);

    // 현재 노드와 부모 노드의 우선순위를 비교한다
    while (parentNode !== 0 && this.compare(this.heap[currentNode], this.heap[parentNode])) {
      // 우선순위가 높다면(낮다면) 부모 노드와 순서를 바꾼다.
      this._swap(parentNode, currentNode);
      // 현재 노드와 부모 노드의 인덱스를 업데이트한다.
      currentNode = parentNode;
      parentNode = Math.floor(currentNode / 2);
    }
  }

  pop() {
    if (this.heap.length === 1) return;
    if (this.heap.length === 2) return this.heap.pop();

    const top = this.heap[1];
    // 루트 노드를 삭제하고 마지막 노드를 루트 노드에 위치시킨다.
    this.heap[1] = this.heap.pop();

    let [currentNode, leftChild, rightChild] = [1, 2, 3];

    // 현재 노드와 자식 노드들의 우선 순위를 비교한다.
    while (
      this.compare(this.heap[leftChild], this.heap[currentNode]) ||
      this.compare(this.heap[rightChild], this.heap[currentNode])
    ) {
      // 자식 노드들 중에서 더 우선순위가 높은(낮은) 노드와 순서를 바꾼다.
      const targetNode = this.compare(this.heap[rightChild], this.heap[leftChild]) ? rightChild : leftChild;
      this._swap(currentNode, targetNode);
      // 현재 노드와 자식 노드들의 인덱스를 업데이트한다.
      currentNode = targetNode;
      leftChild = currentNode * 2;
      rightChild = currentNode * 2 + 1;
    }

    return top;
  }

  size() {
    return this.heap.length - 1;
  }

  top() {
    return this.heap[1];
  }

  _swap(a, b) {
    [this.heap[a], this.heap[b]] = [this.heap[b], this.heap[a]];
  }
}

사용 예시

복사
const minHeap = new PriorityQueue((a, b) => a < b);

minHeap.push(5);
minHeap.push(10);
minHeap.push(-3);
minHeap.push(2);

minHeap.pop(); // -3
minHeap.pop(); // 2
minHeap.pop(); // 5
minHeap.pop(); // 10

const maxHeap = new PriorityQueue((a, b) => a > b);

maxHeap.push(5);
maxHeap.push(10);
maxHeap.push(-3);
maxHeap.push(2);

maxHeap.pop(); // 10
maxHeap.pop(); // 5
maxHeap.pop(); // 2
maxHeap.pop(); // -3