TypeScript로 우선순위 큐를 구현해보자(feat.Jest)

우선순위 큐(Priority Queue)[참고]

우선순위 큐는 큐의 FIFO 구조를 변형한 것으로 가장 우선순위가 높은 데이터가 가장 먼저 나옵니다.

 

배열이나 링크드 리스트, 힙으로 구현할 수 있고, 구현 방법에 따라 아래와 같은 시간 복잡도를 갖습니다.

구현 방식	삽입	제거
unordered array	O(1)	O(n)
unordered linked list	O(1)	O(n)
ordered array	O(n)	O(1)
ordered linked list	O(n)	O(1)
heap	O(log n)	O(log n)

배열이나 링크드 리스트로 구현하는 방법은 매우 간단합니다.

정렬되어 있지 않는 경우, 삽입은 맨 뒤에 하고, 제거할 때 전체를 순회하면서 가장 우선순위가 높은 데이터를 찾아 제거합니다.

정렬되어 있는 경우, 삽입할 때 순회하면서 우선순위에 맞는 위치를 찾아 삽입하고, 제거할 때는 맨 앞에 있는 데이터를 제거합니다.

이 글에서는 Heap으로 구현한 코드를 소개합니다.

구현 내용

• 데이터를 삽입하는 push 구현
• 가장 우선순위가 높은 데이터를 반환하고 제거하는 pop 구현(default: max heap)
• capacity가 부족하거나 많이 남는 경우 resize 하는 기능 구현
• 우선순위 비교 함수를 사용자가 설정할 수 있는 기능 구현
• 가장 우선순위가 높은 데이터를 확인할 수 있는 top, 비어있는지 여부를 확인할 수 있는 empty, 우선순위 큐를 초기화하는 clear 구현

소스코드

PriorityQueue.ts

// PriorityQueue.ts
class PriorityQueue<T> {
  static readonly defaultCapacity = 16;
  static readonly EMPTY_SYMBOL = Symbol('EMPTY');

  private _pq: Array<T>;
  private _capacity: number = PriorityQueue.defaultCapacity;
  private _size: number = 0;
  // a가 b보다 우선순위가 높거나 같을 때 true 반환. default: max heap
  private _compareFn: Function = (a: T, b: T): boolean =>
    typeof a === 'number' ? a >= b : String(a) >= String(b);

  constructor(_compareFn?: Function) {
    _compareFn && (this._compareFn = _compareFn);

    this._pq = new Array(this._capacity + 1);
  }

  get size(): number {
    return this._size;
  }

  push(value: T) {
    if (this._size === this._capacity) {
      this.resize(this._capacity * 2);
    }
    this._pq[++this._size] = value;

    // slide
    let child: number = this._size;
    let parent: number = Math.floor(child / 2);

    while (parent > 0 && this._compareFn(this._pq[child], this._pq[parent])) {
      this.swapByIndex(child, parent);
      child = parent;
      parent = Math.floor(child / 2);
    }
  }

  pop(): T | Symbol {
    if (this.empty()) {
      return PriorityQueue.EMPTY_SYMBOL;
    }
    if (
      this._size > PriorityQueue.defaultCapacity &&
      this._size <= this._capacity / 2
    ) {
      this.resize(Math.floor(this._capacity / 2));
    }
    const returnValue = this._pq[1];

    // slide
    let parent: number = 1;
    this._pq[1] = this._pq[this._size];

    while (true) {
      let leftChild: number = parent * 2;
      let rightChild: number = parent * 2 + 1;

      if (
        rightChild <= this._size &&
        this._compareFn(this._pq[rightChild], this._pq[leftChild]) &&
        this._compareFn(this._pq[rightChild], this._pq[parent])
      ) {
        // 오른쪽 자식과 swap하는 경우
        this.swapByIndex(rightChild, parent);
        parent = rightChild;
      } else if (
        leftChild <= this._size &&
        this._compareFn(this._pq[leftChild], this._pq[parent])
      ) {
        // 왼쪽 자식과 swap하는 경우
        this.swapByIndex(leftChild, parent);
        parent = leftChild;
      } else {
        // swap을 모두 마친 경우
        break;
      }
    }

    this._size--;
    return returnValue;
  }

  top(): T | Symbol {
    return this.empty() ? PriorityQueue.EMPTY_SYMBOL : this._pq[1];
  }

  empty(): boolean {
    return this._size === 0;
  }

  clear() {
    this._size = 0;
    this._capacity = PriorityQueue.defaultCapacity;
    this._pq = new Array(this._capacity + 1);
  }

  private resize(newCapacity: number) {
    const newPQ = new Array(newCapacity + 1);
    this._capacity = newCapacity;

    for (let i = 1; i <= this.size; ++i) {
      newPQ[i] = this._pq[i];
    }
    this._pq = newPQ;
  }

  private swapByIndex(indexA: number, indexB: number) {
    let tmp: T = this._pq[indexA];
    this._pq[indexA] = this._pq[indexB];
    this._pq[indexB] = tmp;
  }
}

export default PriorityQueue;

PriorityQueue.test.ts

// PriorityQueue.test.ts
import PriorityQueue from './PriorityQueue';

describe('priority queue test', () => {
  let pq: PriorityQueue<number>;

  beforeEach(() => {
    pq = new PriorityQueue<number>();
  });

  it('new 연산자로 PriorityQueue 인스턴스를 생성한다.', () => {
    expect(Object.getPrototypeOf(pq)).toEqual(PriorityQueue.prototype);
  });

  it('push - 값을 추가한다.', () => {
    pq.push(1);
    expect(pq.size).toBe(1);

    pq.push(2);
    expect(pq.size).toBe(2);
  });

  it('push - 꽉 찬 상태에서 값을 추가하면 resize 후, 값을 추가한다.', () => {
    for (let i = 0; i <= PriorityQueue.defaultCapacity; ++i) {
      pq.push(i);
    }
    expect(pq.size).toBe(PriorityQueue.defaultCapacity + 1);
  });

  it('pop - 비어있으면 Symbol을 반환한다.', () => {
    expect(pq.pop()).toBe(PriorityQueue.EMPTY_SYMBOL);
  });

  it('pop - 우선순위가 가장 높은 값을 제거하고 반환한다(default: max heap).', () => {
    pq.push(1);
    pq.push(1);
    pq.push(2);
    pq.push(3);

    expect(pq.pop()).toBe(3);
    expect(pq.pop()).toBe(2);
    expect(pq.pop()).toBe(1);
    expect(pq.pop()).toBe(1);
    expect(pq.size).toBe(0);
  });

  it('pop - capacity와 size가 2배 이상 차이나면 resize한다.', () => {
    for (let i = 0; i <= PriorityQueue.defaultCapacity * 2; ++i) {
      pq.push(i);
    }
    for (let i = 0; i <= PriorityQueue.defaultCapacity; ++i) {
      pq.pop();
    }
  });

  it('top - 비어있으면 Symbol을 반환한다.', () => {
    expect(pq.top()).toBe(PriorityQueue.EMPTY_SYMBOL);
  });

  it('top - 다음에 pop될 값(가장 우선순위가 높은 값)을 반환한다.', () => {
    pq.push(1);
    expect(pq.top()).toBe(1);

    pq.push(3);
    expect(pq.top()).toBe(3);

    pq.push(2);
    expect(pq.top()).toBe(3);
  });

  it('empty - 비어있으면 true, 차있으면 false를 반환한다.', () => {
    expect(pq.empty()).toBe(true);

    pq.push(1);
    expect(pq.empty()).toBe(false);
  });

  it('clear - PriorityQueue를 초기화한다.', () => {
    pq.push(1);
    pq.push(2);
    pq.clear();

    expect(pq.empty()).toBe(true);
  });

  it('compareFn - 숫자가 아닌 값의 경우에는 String으로 변환하여 우선순위를 비교한다(string).', () => {
    const strPQ = new PriorityQueue<string>();

    strPQ.push('a');
    strPQ.push('b');
    strPQ.push('c');
    strPQ.push('d');

    expect(strPQ.pop()).toBe('d');
    expect(strPQ.pop()).toBe('c');
    expect(strPQ.pop()).toBe('b');
    expect(strPQ.pop()).toBe('a');
  });

  it('compareFn - 숫자가 아닌 값의 경우에는 String으로 변환하여 우선순위를 비교한다(boolean).', () => {
    const boolPQ = new PriorityQueue<boolean>();

    boolPQ.push(false);
    boolPQ.push(true);

    expect(boolPQ.pop()).toBe(true);
    expect(boolPQ.pop()).toBe(false);
  });

  it('compareFn - 우선순위 비교 함수를 설정할 수 있다(min heap).', () => {
    const compareFn = (a: number, b: number): boolean => b > a;
    const minHeap = new PriorityQueue<number>(compareFn);

    minHeap.push(1);
    minHeap.push(2);
    minHeap.push(3);
    minHeap.push(4);

    expect(minHeap.pop()).toBe(1);
    expect(minHeap.pop()).toBe(2);
    expect(minHeap.pop()).toBe(3);
    expect(minHeap.pop()).toBe(4);
  });

  it('compareFn - 우선순위 비교 함수를 설정할 수 있다(object).', () => {
    type Person = {
      name: string;
      age: number;
    };
    const compareFn = (a: Person, b: Person): boolean => a.age > b.age;
    const customPQ = new PriorityQueue<Person>(compareFn);

    customPQ.push({ name: 'choi', age: 10 });
    customPQ.push({ name: 'shim', age: 20 });
    customPQ.push({ name: 'lee', age: 30 });
    customPQ.push({ name: 'shin', age: 40 });
    customPQ.push({ name: 'park', age: 50 });

    expect(customPQ.pop()).toEqual({ name: 'park', age: 50 });
    expect(customPQ.pop()).toEqual({ name: 'shin', age: 40 });
    expect(customPQ.pop()).toEqual({ name: 'lee', age: 30 });
    expect(customPQ.pop()).toEqual({ name: 'shim', age: 20 });
    expect(customPQ.pop()).toEqual({ name: 'choi', age: 10 });
  });
});

기타

코드 내 연산을 간단하게 만들기 위해 배열의 0번째 인덱스는 비워두고 1번째 인덱스부터 사용했습니다.