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

이 글을 쓰게된 계기

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

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

힙(Heap)이란?

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

• peekMin, peekMax : $O(1)$
• insert : $O(logN)$
• remove : $O(logN)$

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

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

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

구현 (최소 힙)

힙의 구조는 크게 아래와 같이 나눌 수 있다.

• 삽입
• 힙의 조건에 맞게 데이터 재배치하기 (heap.length-1 -> 0)
• 삭제
• 힙의 조건에 맞게 데이터 재배치하기 (0 -> heap.length-1)

그리고 부모 노드, 왼쪽 자식노드, 오른쪽 자식노드의 인덱스를 얻기 위한 식은 아래와 같다.

• parentIndex = Math.floor((index - 1) / 2)
• leftChildIndex = index * 2 + 1
• rightChildIndex = index * 2 + 2

그리고 이를 아래처럼 비트 시프트 연산 방식으로 계산하는 방법도 있다.

• parentIndex = (index - 1) >> 1
• leftChildIndex = (index << 1) + 1
• rightChildIndex = (index << 1) + 2

비트 시프트 연산을 사용하면 연산 속도가 더 빠르고, 굳이 Math.floor 함수를 사용하지 않아도 나머지가 떨어져 나가기 때문에 비트 시프트 연산을 더 선호하는 편이다. 이를 이용해 아래와 같이 간결하게 함수를 만들어줄 수도 있다.

getLeftChildIndex = (index) => (index << 1) + 1;
getRightChildIndex = (index) => (index << 1) + 2;
getParentIndex = (index) => (index - 1) >> 1;

하지만 막상 힙을 구현할 때 위 함수들의 재사용성이 낮기 때문에 나는 함수를 구현해주지 않고, 필요할 때 직접 인덱스를 구하는 방식으로 구현했다.

생성자

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

삽입

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

조건에 맞게 재배치 (삽입된 원소 위로 올리기)

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

insert 함수에서 넘겨준 heap 의 마지막 인덱스를 이용해서 인덱스가 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;
}

min 변수에 루트 노드를 저장해놓고, end 변수에는 heap 에 마지막 노드를 꺼낸 뒤 할당해준다.
힙에 남은 원소가 없다면(0 이라면) 그냥 min 을 리턴해주고 아니라면 루트노드(heap[0])에 end 를 할당해주고 heapifyDown을 수행하여 재배치를 한 뒤 min 을 리턴해준다.

조건에 맞게 재배치 (루트노드를 아래로 끌어 내리기)

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

remove 함수에서 0번 인덱스를 인자로 전달받아서 왼쪽 자식 노드가 있다면 왼쪽 자식노드와 현재 노드를 비교하고, 오른쪽 자식노드가 있다면 오른쪽 자식노드랑도 비교하여 만약 자식 노드의 값이 더 작다면 더 작은 자식 노드의 위치와 현재 노드의 위치를 바꿔준다. 만약 자식 노드중 더 작은 노드가 없다면 조건에 맞게 배치된 것이므로 반복문을 탈출한다.

완성

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

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

    heapifyUp(index) {
        while (index > 0) {
            const parentIndex = (index - 1) >> 1;
            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 << 1) + 1;
            const right = (index << 1) + 2;
            let smallest = index;
            if (this.heap[left] && this.heap[left] < this.heap[smallest]) {
                smallest = left;
            }
            if (this.heap[right] && this.heap[right] < this.heap[smallest]) {
                smallest = right;
            }
            if (smallest === index) break;
            [this.heap[index], this.heap[smallest]] = [this.heap[smallest], this.heap[index]];
            index = smallest;
        }
    }

    isEmpty() {
        return this.heap.length === 0;
    }
}

최대힙을 구현하고 싶다면 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 = (index - 1) >> 1;
            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 << 1) + 1;
            const right = (index << 1) + 2;
            let smallest = index;
            if (this.heap[left] && this.heap[left].priority < this.heap[smallest].priority) {
                smallest = left;
            }
            if (this.heap[right] && this.heap[right].priority < this.heap[smallest].priority) {
                smallest = right;
            }
            if (smallest === index) break;
            [this.heap[index], this.heap[smallest]] = [this.heap[smallest], this.heap[index]];
            index = smallest;
        }
    }

    isEmpty() {
        return this.heap.length === 0;
    }
}

참고

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