【TypeScript】優先度付きキュー

概要

優先順位の高い作業から順番にストリーミング処理したいことってありますよね。
もちろん、TypeScriptのListを使えば簡単に実装できます。
しかし、毎回Listの中身を一件ずつ見ていき一番優先度の高い要素を探しだして取り出すのは効率が悪いです。
そこで、それを高速に扱えるデータ構造を紹介します。

要件

1. 要素を追加できる
2. 今まで追加した要素の中で一番優先度が大きい値を答えられる
3. 一番大きい要素を取り出す（削除する）ことができる
4. 3の時、二番目に大きい要素が一番に昇格する

これらを高速に処理できればいいです。

コード

ひとまず要素が整数である場合に限って考えます。

class PriorityQueue {
  private list = new Array<number>();

  public push(v: number) {
    this.list.push(v);
    const back = this.list.length - 1;

    this.against(back);
  }

  private against(child: number) {
    if (child === 0) return;
    const parent = Math.ceil(child / 2) - 1;

    const childValue = this.list[child];
    const parentValue = this.list[parent];

    if (parentValue < childValue) {
      this.swap(child, parent);
      this.against(parent);
    }
  }

  private swap(a: number, b: number) {
    const tmp = this.list[a];
    this.list[a] = this.list[b];
    this.list[b] = tmp;
  }

  public top(): number {
    if (this.list.length === 0) throw new Error("empty");
    return this.list[0];
  }

  public pop(): number {
    if (this.list.length === 0) throw new Error("empty");
    const ans = this.top();
    this.list[0] = this.list.pop();
    this.flow(0);
    return ans;
  }

  private flow(parent: number) {
    const parentValue = this.list[parent];

    const left = parent * 2 + 1;
    const right = parent * 2 + 2;

    if (!this.inRange(left)) return;

    // 左子入るが右子はいない
    if (this.inRange(left) && !this.inRange(right)) {
      const leftValue = this.list[left];

      if (parentValue < leftValue) {
        this.swap(parent, left);
      }
    }

    if (this.inRange(left) && this.inRange(right)) {
      const leftValue = this.list[left];
      const rightValue = this.list[right];
      const target = leftValue > rightValue ? left : right;
      const targetValue = this.list[target];

      if (parentValue < targetValue) {
        this.swap(parent, target);
        this.flow(target);
      }
    }
  }

  private inRange(index: number): boolean {
    return index < this.list.length;
  }

  public size() {
    return this.list.length;
  }
}

function main() {
  const pq = new PriorityQueue();
  pq.push(8);
  pq.push(2);
  console.log(pq.pop());
  pq.push(10);
  console.log(pq.pop());
  pq.push(11);
  console.log(pq.pop());
  console.log(pq.pop());
}

解説

二分ヒープというデータ構造を応用しています。
二分ヒープは木構造です。
この木構造のルールはたった一つです。

親の値は必ず子の値以上になる
兄弟に大小関係はありません。

すると、木のてっぺんが常に最も大きい値になります。

要素の追加の仕方

1. 木の一番下に要素を追加する
2. 『親の値は必ず子の値以上になる』 を満たすまで上に遡ってswapをしていく

要素の取り出しの仕方

1. 木の一番上の値を答える
2. 木の一番上に一番下の値を移す
3. 木の長さを一個縮める
4. 『親の値は必ず子の値以上になる』 を満たすまで上から下にswapをしていく

4の時、二人の子の中で大きい方を選んでswapをしていくと色々都合がいいらしいです。

簡単ですね。

計算量

各要素が2個ずつ子供を持つので、
木の高さはlog2(要素数)になります。
なので、要素の追加と取り出しの最悪計算量はlog2(要素数)です。

拡張

今回は常に一番大きい値を返すキューを作りました。
ジェネリクスと高階関数を使って色んな大小関係を扱えるようにすると便利になりますね。