😽

【TypeScript】最短距離を求める『ダイクストラ』

2021/03/29に公開

やりたいこと

こういう国があったとします。

街0から街1~5までの最短距離を求めたいとします。

このようなときにはダイクストラ法というアルゴリズムを使うと便利です。

コード

// aの優先度が高ければtrue
type Compare<T> = (a: T, b: T) => boolean;

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

  constructor(private compare: Compare<T>) {}

  public push(v: T) {
    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 (this.compare(childValue, parentValue)) {
      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(): T {
    if (this.list.length === 0) throw new Error("empty");
    return this.list[0];
  }

  public pop(): T {
    if (this.list.length === 0) throw new Error("empty");

    if (this.list.length === 1) {
      const ans = this.top();
      this.list.pop();
      return ans;
    }

    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(right)) {
      // 左子はいるが右子はいない
      const leftValue = this.list[left];

      if (this.compare(leftValue, parentValue)) {
        this.swap(parent, left);
      }
      return;
    }
    const leftValue = this.list[left];
    const rightValue = this.list[right];
    const target = this.compare(leftValue, rightValue) ? left : right;
    const targetValue = this.list[target];

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

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

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

interface Task {
  checkPoint: number;
  totalCost: number;
}

interface Edge {
  from: number;
  to: number;
  cost: number;
}

class Dijkstra {
  edges: Edge[][];

  public constructor(private n: number) {
    this.edges = new Array(n);
    for (let i = 0; i < n; i++) this.edges[i] = [];
  }

  public insert(from: number, to: number, cost: number): void {
    this.edges[from].push({ from, to, cost });
  }

  public build(start: number): number[] {
    const distances = new Array<number>(this.n);
    for (let i = 0; i < this.n; i++) distances[i] = Infinity;

    distances[start] = 0;

    const queue = new PriorityQueue<Task>((a, b) => a.totalCost < b.totalCost);

    queue.push({ checkPoint: start, totalCost: 0 });

    while (queue.size() > 0) {
      const t = queue.pop();

      if (t.totalCost > distances[t.checkPoint]) continue;

      for (const e of this.edges[t.checkPoint]) {
        if (distances[e.from] + e.cost < distances[e.to]) {
          distances[e.to] = distances[e.from] + e.cost;
          queue.push({ checkPoint: e.to, totalCost: distances[e.to] });
        }
      }
    }
    return distances;
  }
}

function main() {
  const d = new Dijkstra(6);
  d.insert(0, 1, 2);
  d.insert(0, 3, 4);
  d.insert(1, 3, 1);
  d.insert(2, 5, 4);
  d.insert(3, 2, 3);
  d.insert(3, 4, 5);
  d.insert(4, 5, 1);
  const ans = d.build(0);
  console.log(ans);
}

riku@rikuubuntu:~/algorithm$ node dk.js
[ 0, 2, 6, 3, 8, 9 ]

解説

前回解説した優先度付きキューを使います。
スタート地点から始めて、短い道から順番に処理していきます。

計算量

多くの場合、街の数より道の数のほうが多いと思います。
優先度付きキューを使って各道を一回ずつ処理するので、大体O(\log(道の数))です。

Discussion