Я реализовал класс , чтобы сделать именно то, что вам нужно:
- вставка
enqueue
- выдержка в мин
dequeue
- Ключ уменьшения -
putValue
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
}
}