자바스크립트로 힙(Heap) 구현하기

이 글을 쓰게된 계기

최근 다익스트라 알고리즘 문제를 처음 풀어 봤는데, 다익스트라 알고리즘에서 동작 속도를 빠르게 하려면 우선순위 큐(Priority Queue)를 사용해야 했다. 하지만 자바스크립트와 Node.js에서는 힙을 구현해주는 내장 모듈이 없기 때문에 힙을 직접 구현해서 문제를 풀어야 했다.

코딩테스트에서 이걸 직접 구현하려면 코드를 완벽히 숙지해두어야 하고, 따라서 코드를 정리해서 외우기 편하게 만들어 놓으면 좋을것 같아서 이 글을 작성하게 되었다.

힙(Heap)이란?

힙은 트리 기반의 자료구조이고 데이터 탐색, 삽입, 삭제의 시간 복잡도는 아래와 같다.

• 최솟값 탐색, 최댓값 탐색 : $O(1)$
• 삽입 : $O(logN)$
• 삭제 : $O(logN)$

힙은 데이터를 저장하는 방식에 따라 크게 두 가지로 나눌 수 있다.

• 최대 힙(Max Heap)
  부모 노드는 항상 자식 노드보다 크거나 같다.
• 최소 힙(Min Heap)
  부모 노드는 항상 자식 노드보다 작거나 같다.

힙은 '완전 이진 트리' 구조를 사용하는데, 이진 트리란 하나의 부모 노드는 두 개의 자식 노드만 가질 수 있음을 의미하고, 여기서 '완전' 이진 트리란 트리의 가장 아래 층을 제외한 모든 층에 노드가 완전히 채워져 있어야 함을 의미한다.

구현 (최소 힙)

힙의 동작은 크게 아래와 같이 나눌 수 있다.

• 삽입
• 삭제
• 힙의 조건에 맞게 노드 재배치하기

위의 힙의 시간복잡도를 살펴보면 삽입, 삭제 의 시간복잡도가 $O(logN)$ 인 것을 볼 수 있다. 이는 삽입 과 삭제된 노드의 값에 따라 트리를 업데이트해줘야 되기 때문이다.

생성자

constructor() {
  	this.heap = [];
}

삽입 & 노드 재배치

insert(value) {
    this.heap.push(value);
    this.heapifyUp(this.heap.length - 1); 
}

heapifyUp(index) {
    while (index > 0) {
      const parentIndex = Math.floor((index - 1) / 2);
      if (this.heap[parentIndex] <= this.heap[index]) break;
      [this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
      index = parentIndex;
    }
}

노드를 트리에 push하고, heapifyUp 을 실행시켜 새로 추가된 노드와 해당 노드의 부모 노드의 값을 비교하여 값이 더 작은 노드가 부모 노드가 되도록 위치를 조정해준다.

이를 루트 노드(index=0)까지 반복하여 '최소힙'을 보장한다. 이진 트리이므로 이 과정의 시간 복잡도는 $O(logN)$ 이다.

삭제 & 노드 재배치

remove() {
    const min = this.heap[0];
    const end = this.heap.pop();
    if (this.heap.length > 0) {
      this.heap[0] = end;
      this.heapifyDown(0);
    }
    return min;
}

heapifyDown(index) {
    while (index < this.heap.length) {
      const left = index * 2 + 1;
      const right = index * 2 + 2;
      let smaller = index;
      if (this.heap[left] && this.heap[left] < this.heap[smaller]) {
        smaller = left;
      }
      if (this.heap[right] && this.heap[right] < this.heap[smaller]) {
        smaller = right;
      }
      if (smaller === index) break;
      [this.heap[index], this.heap[smaller]] = [this.heap[smaller], this.heap[index]];
      index = smaller;
    }
}

루트 노드가 최솟값이므로 min 변수에 루트 노드를 저장해놓고 마지막 인덱스의 노드(end)를 꺼낸다.
노드를 꺼낸 후, 힙에 남은 원소가 없다면 그냥 min === end 이므로 둘 중 무엇을 리턴해줘도 상관없다. 힙에 남은 원소가 있다면 루트 노드(heap[0])에 end를 할당하고 노드를 재배치한다(heapifyDown).

heapifyDown은 루트 노드부터 아래로 내려가며 자식 노드와 부모 노드를 비교하여 '최소힙'을 보장하도록 하는 로직이다. 이진 트리이기 때문에 왼쪽 자식 노드, 오른쪽 자식 노드와 비교하여 자식 노드가 더 작다면 자식 노드의 위치와 부모 노드의 위치를 바꾼다.

완성

class MinHeap {
    constructor() {
        this.heap = [];
    }

    insert(value) {
        this.heap.push(value);
        this.heapifyUp(this.heap.length - 1); 
    }

    heapifyUp(index) {
        while (index > 0) {
          const parentIndex = Math.floor((index - 1) / 2);
          if (this.heap[parentIndex] <= this.heap[index]) break;
          [this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
          index = parentIndex;
        }
    }
  
    remove() {
        const min = this.heap[0];
        const end = this.heap.pop();
        if (this.heap.length > 0) {
          this.heap[0] = end;
          this.heapifyDown(0);
        }
        return min;
    }

    heapifyDown(index) {
        while (index < this.heap.length) {
          const left = index * 2 + 1;
          const right = index * 2 + 2;
          let smaller = index;
          if (this.heap[left] && this.heap[left] < this.heap[smaller]) {
            smaller = left;
          }
          if (this.heap[right] && this.heap[right] < this.heap[smaller]) {
            smaller = right;
          }
          if (smaller === index) break;
          [this.heap[index], this.heap[smaller]] = [this.heap[smaller], this.heap[index]];
          index = smaller;
        }
    }
}

최대힙을 구현하고 싶다면 heapifyUp 과 heapifyDown 에서 조건을 반대로 적용해주면 된다.

우선순위 큐

최소힙을 이용하여 아래와 같이 우선순위 큐를 구현할 수 있다.

class PriorityQueue {
    constructor() {
        this.heap = [];
    }

    enqueue(node, priority) {
        this.heap.push({ node, priority });
        this.heapifyUp(this.heap.length - 1);
    }

    heapifyUp(index) {
        while (index > 0) {
            const parentIndex = Math.floor((index - 1) / 2);
            if (this.heap[parentIndex].priority <= this.heap[index].priority) break;
            [this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
            index = parentIndex;
        }
    }

    dequeue() {
        const min = this.heap[0];
        const end = this.heap.pop();
        if (this.heap.length > 0) {
            this.heap[0] = end;
            this.heapifyDown(0);
        }
        return min;
    }

    heapifyDown(index) {
        while (index < this.heap.length) {
            const left = index * 2 + 1;
            const right = index * 2 + 2;
            let smaller = index;
            if (this.heap[left] && this.heap[left].priority < this.heap[smaller].priority) {
                smaller = left;
            }
            if (this.heap[right] && this.heap[right].priority < this.heap[smaller].priority) {
                smaller = right;
            }
            if (smaller === index) break;
            [this.heap[index], this.heap[smaller]] = [this.heap[smaller], this.heap[index]];
            index = smaller;
        }
    }
}

참고

Geeks for Geeks Min Heap
JavaScript로 Heap | Priority Queue 구현하기