Scalaでダイクストラをやってみる

経緯

今度新しくやるプロジェクトはScalaでやろうと思っており、練習がてらにScalaでダイクストラ法を実装してみました。まだScalaをやり始めて3日なので、関数型っぽくやろうとおもったのですが、たんにScalaの強力なコレクションに頼っただけみたいになりました。

今回はデータ構造は簡単に実装するだけなので隣接行列で用意しました。今からScala始めようと思っている人がこれを見て、for文とかなしで実装できるんだなとか思ってくれたらうれしいです笑

コード

import scala.annotation.tailrec

object Dijkstra {
  def main(args: Array[String]) {
    val distance = List(
      List(0, 1, 2, 0, 0, 0),
      List(1, 0, 1, 2, 0, 0),
      List(2, 1, 0, 0, 1, 0),
      List(0, 2, 0, 0, 2, 2),
      List(0, 0, 1, 2, 0, 1),
      List(0, 0, 0, 2, 1, 0)
    )

    val init = List(0, -1, -1, -1, -1, -1)

    def getMinDistance(list: List[Int])(i: Int): Int =
      (0 to list.size - 1).toList
        .filter(list(_) != -1)
        .filter(distance(_)(i) != 0)
        .map(j => distance(i)(j) + list(j))
        .min

    def isIsolated(list: List[Int])(i: Int): Boolean =
      (0 to list.size - 1).toList
        .filter(list(_) != -1)
        .filter(distance(_)(i) != 0)
        .isEmpty

    def getNextMinPosition(list: List[Int]): Int =
      (0 to list.size - 1).toList
        .filter(list(_) == -1)
        .filterNot(isIsolated(list))
        .minBy(getMinDistance(list))

    def compute(l: List[Int]): List[Int] = {
      @tailrec
      def rec(list: List[Int], restCount: Int): List[Int] = {
        restCount match {
          case 0 => list
          case _ => {
            val next = getNextMinPosition(list)
            val dis = getMinDistance(list)(next)
            rec(list.updated(next, dis), restCount - 1)
          }
        }
      }
      rec(l, l.size - 1)
    }

    println(compute(init))
  }
}

       2
    b-----d
 1 /|     |\2
  a |1   2| f
 2 \|     |/1
    c-----e
       1

隣接行列はこんな感じになっています。C言語とかで実装するとループで配列を処理していきますが、Scalaには強力なコレクションが用意されているので見てすぐわかるくらい簡潔に書けました。getNextMinPositionが二か所で呼ばれているのはちょっとダサいですが、Scala初心者なので多めに見てください。