Scalaで優先度付きキュー

30分プログラム、その511。Scalaで優先度付きキューを作ってみよう。
wikipedia:優先度つきキューによると、連想リストとキューを組合せるのが一番楽らしいので、それをそのままScalaに落し込んでやる。

ちなみに、Scala Standard Library 2.12.8 - scala.collection.mutable.PriorityQueueというクラスがあるのは知っている。

使い方

val q1 = PriorityQueue[Int]()
val q2 = q1.enqueue(0,1).enqueue(0,2).enqueue(-1,3)

// 3が一番、優先度が高い
val (x,q3) = q2.dequeue // x = 3

// あとは入れた順にでてくる
val (y,q4) = q3.dequeue // y = 1
val (z,q5) = q4.dequeue // z = 2

ソースコード

import scala.collection.immutable._

class PriorityQueue[T](map : Map[Int,Queue[T]]){
  def enqueue(n : Int,v : T) : PriorityQueue[T] = {
    val queue  = map(n).enqueue(v)
    val newMap = map + (n -> queue)
    new PriorityQueue(newMap)
  }

  def dequeue() : (T,PriorityQueue[T]) = {
    val n = map.keys.toList.sort(_ < _).head
    val (value,queue) = map(n).dequeue
    val newMap = if(queue.isEmpty){
      map - n
    }else{
      map + (n -> queue)
    }
    (value,new PriorityQueue(newMap))
  }
}

object PriorityQueue {
  def apply[T]() : PriorityQueue[T] = {
    val empty = new Queue[T]
    new PriorityQueue[T](Map[Int,Queue[T]]() withDefaultValue empty)
  }
}

val q1 = PriorityQueue[Int]()
val q2 = q1.enqueue(0,1).enqueue(0,2).enqueue(-1,3)

val (x,q3) = q2.dequeue // x = 3
val (y,q4) = q3.dequeue // y = 1
val (z,q5) = q4.dequeue // z = 2