An Mutable PriorityQueue that allows random acces to elements and allows changins values

heapedMap.scala

import scala.annotation.tailrec
import scala.collection.mutable.{ArrayBuffer, Map}
import scala.util.Try

class HeapedMap[K, V] (implicit ord: Ordering[V]){
  val dict = Map[K,Int]() // Keeps index of key withing vector
  val vector = ArrayBuffer[(K,V)]()

  override def toString(): String = vector.toString
  def toMap(): scala.collection.immutable.Map[K,V] = dict.mapValues(vector(_)._2).toMap
  def toSeq(): Seq[(K,V)] = vector
  def toList(): List[(K,V)] = vector.toList

  private def parent(i: Int): Int = (i - 1) / 2
  private def right(i: Int): Int = (i * 2 + 1) + 1
  private def left(i: Int): Int = (i * 2 + 1)
  private def exists(i: Int): Boolean = Try(vector(i)).isSuccess
  private def compare(x: V, y: V): Boolean = ord.lteq(x,y)
  private def swap(i: Int, j: Int): Unit = {
    val aux = vector(i)
    vector(i) = vector(j)
    dict(vector(i)._1) = i
    vector(j) = aux
    dict(vector(j)._1) = j
  }
  private def replace(i: Int, j: Int): Unit = {
    dict -= vector(i)._1
    vector(i) = vector(j)
    dict(vector(i)._1) = i
    vector.remove(j)
  }

  private def insert(key: K, value: V): Unit = {
    vector += ((key, value))
    dict(key) = vector.size - 1
    bubbleUp(vector.size - 1)
  }

  private def delete(i: Int): Unit = {
    if (vector.size == 1) {
      dict -= vector(0)._1
      vector.remove(0)
    } else {
      replace(i, vector.size - 1)
      bubbleDown(i)
    }
  }

  def isEmpty(): Boolean = vector.isEmpty

  def enqueue(key: K, value: V): Unit = {
    if (!dict.contains(key)) {
      insert(key,value)
    }
    // TODO: handle when it already contains the key
  }

  @tailrec
  private def bubbleUp(i: Int): Unit = {
    val p = parent(i)
    if ((p != i) && (!compare(vector(p)._2, vector(i)._2))) {
      swap(p, i)
      bubbleUp(p)
    }
  }

  @tailrec
  private def bubbleDown(i: Int): Unit = {
    var largest = i
    val l = left(i)
    val r = right(i)
    if ((exists(l)) && (compare(vector(l)._2, vector(largest)._2))) {
      largest = l
    }
    if ((exists(r)) && (compare(vector(r)._2, vector(largest)._2))) {
      largest = r
    }

    if (largest != i) {
      swap(i, largest)
      bubbleDown(largest)
    }
  }

  def dequeue(): Option[(K, V)] = {
    val optionRoot = vector.headOption
    if (optionRoot.isDefined) {
        delete(0)
    }
    optionRoot
  }

  def dequeueAll(): Seq[(K,V)] = {
    val resp = ArrayBuffer[(K,V)]()
    var x = dequeue
    while (x.isDefined) {
      resp += x.get
      x = dequeue
    }
    resp
  }

  def headOption(): Option[(K,V)] = vector.headOption
  def get(k: K): Option[V] = {
    dict.get(k) match {
      case Some(i) => Some(vector(i)._2)
      case None => None
    }
  }

  // change values and heapify
  // * PriorityQueue does not have this feature
  def putValue(key: K, value: V): Unit = {
    val optionOldValue = get(key)
    if (optionOldValue.isDefined) {
      val oldValue = optionOldValue.get
      val i = dict(key)
      vector(i) = (key, value)
      if (compare(value, oldValue)) {
        bubbleUp(i)
      } else {
        bubbleDown(i)
      }
    } else {
      // if key does not exist, insert it
      insert(key,value)
    }
  }

  // different from dequeue, removes an arbitrary element
  def remove(key: K): Unit = {
    if (dict.contains(key)) {
      delete(dict(key))
    }
    // TODO: handle when it does not contain the key
  }
}


// -------------
// Usage Example
// -------------

case class CC(key: Int, value: Double) extends Ordered[CC] {
  def compare(that: CC) = this.value.compareTo(that.value)
}

def Desc[T : Ordering] = implicitly[Ordering[T]].reverse

val heap = new HeapedMap[Int, CC]()(Desc)

println("Enqueue")

for (i <- 4 to 20) {
  heap.enqueue(i, CC(i,i * 5.0))
}

println(heap)

heap.enqueue(2, CC(2, 10.0))
heap.enqueue(3, CC(3, 15.0))
heap.enqueue(1, CC(1, 5.0))
heap.enqueue(21, CC(21, 21*5.0))

println(heap)

println
println("Dequeue")

println(heap.dequeue)
println(heap)

println(heap.dequeue)
println(heap)

println(heap.dequeue)
println(heap)

println
println("HeadOption")

println(heap.headOption)

println
println("Get")

println(10, heap.get(10))
println(7, heap.get(7))
println(19, heap.get(19))

println
println("PutValue")

heap.putValue(10, CC(10, 500.0))
println(10, heap.get(10))
println(heap)

heap.putValue(13, CC(13, 1.0))
println(13, heap.get(13))
println(heap)

heap.putValue(30, CC(30, 30*5.0))
println(30, heap.get(30))
println(heap)

println
println("Remove")
heap.remove(18)
println(heap)

println
println("DequeueAll")
println(heap.dequeueAll)