Scalaで競技プログラミング: ダイクストラ法

C++で書いたライブラリをScalaで書き直しています。ダイクストラ法全体として以下のようなコードになりました。

case class Edge(to: Int, w: Long)

case class Graph(n: Int) {
  val g: Array[Array[Edge]] = Array.fill(n)(Array.empty)

  def push(from: Int, edge: Edge): Unit = {
    g(from) :+= edge
  }

  def dijkstra(start: Int, Inf: Long = 1e18.toLong): Array[Long] = {
    val d = Array.fill(n)(Inf)
    d(start) = 0
    val queue = mutable.PriorityQueue.empty[(Long, Int)](Ordering.Tuple2[Long, Int].reverse)
    queue.enqueue((0, start))

    while (queue.nonEmpty) {
      val (minCost, minV) = queue.dequeue()
      if (d(minV) >= minCost) {
        g(minV).foreach { e =>
          if (d(e.to) > d(minV) + e.w) {
            d(e.to) = d(minV) + e.w
            queue.enqueue((d(e.to), e.to))
          }
        }
      }
    }
    d
  }
}

今回の主な新情報は PriorityQueue についての以下の行です。

 val queue = mutable.PriorityQueue.empty[(Long, Int)](Ordering.Tuple2[Long, Int].reverse)

Scala Standard Library 2.12.0 - scala.collection.mutable.PriorityQueue

元々Scalaの PriorityQueue は降順で管理されています。今回は、(コスト: Long, 頂点id: Int） というタプルを PriorityQueue で管理します。その際は以下のように順序が規定されています。最初に第一要素を比べ、同じだったら第二要素を比べます。

def compare(x: (T1, T2), y: (T1, T2)): Int = {
  val compare1 = ord1.compare(x._1, y._1)
  if (compare1 != 0) return compare1
  ord2.compare(x._2, y._2)
}

ダイクストラで用いる際には、ある頂点までのコストが軽い順に並べ替えられていた方が嬉しいので、昇順になっている必要があります。そこで、 Ordering を引数に与えることで、並べ替えの順序を指定してあげます。降順の逆の昇順に並べたいので、Ordering.Tuple2[Long, Int].reverse を指定することになります。

他にも dequeue がpopするだけでなく要素も返すところは C++ と異なりますね。（むしろ返す方が自然だとは思いますが...）

AOJの問題でverifyしました。入出力含めた全てのコードはこちらをご覧ください。

github.com