Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import scala.collection.mutable
/**
* Bounded priority queue trait that is intended to be mixed into instances of
* scala.collection.mutable.PriorityQueue. By default PriorityQueue instances in
* Scala are unbounded. This trait modifies the original PriorityQueue's
* enqueue methods such that we only retain the top K elements.
* The top K elements are defined by an implicit Ordering[A].
* @author Ryan LeCompte (lecompte@gmail.com)
*/
trait BoundedPriorityQueue[A] extends mutable.PriorityQueue[A] {
def maxSize: Int
override def +=(a: A): this.type = {
if (size < maxSize) super.+=(a)
else maybeReplaceLowest(a)
this
}
override def ++=(xs: TraversableOnce[A]): this.type = {
xs.foreach { this += _ }
this
}
override def +=(elem1: A, elem2: A, elems: A*): this.type = {
this += elem1 += elem2 ++= elems
}
override def enqueue(elems: A*) {
this ++= elems
}
private def maybeReplaceLowest(a: A) {
// note: we use lt instead of gt here because the original
// ordering used by this trait is reversed
if (ord.lt(a, head)) {
dequeue()
super.+=(a)
}
}
}
object BoundedPriorityQueue {
/**
* Creates a new BoundedPriorityQueue instance.
* @param maximum the max number of elements
* @return a new bounded priority queue instance
*/
def apply[A: Ordering](maximum: Int): BoundedPriorityQueue[A] = {
// note: we reverse the ordering here because the mutable.PriorityQueue
// class uses the highest element for its head/dequeue operations.
val ordering = implicitly[Ordering[A]].reverse
new mutable.PriorityQueue[A]()(ordering) with BoundedPriorityQueue[A] {
implicit override val ord = ordering
override val maxSize = maximum
}
}
}
///////////////// SAMPLE USAGE ////////////////////
scala> val q = BoundedPriorityQueue[Int](5)
q: BoundedPriorityQueue[Int] = PriorityQueue()
scala> val nums = scala.util.Random.shuffle(Vector.range(0, 100000))
nums: scala.collection.immutable.Vector[Int] = Vector(62499, 20240, 71494, 67205, 43994, 31339, 14881, 23726, 54180, 22249, 34191, 8200, 78496, 18751, 86506, 15764, 90615, 61182, 52616, 54248, 48928, 40625, 35891, 72519, 6627, 79008, 40216, 37938, 28009, 59632, 22050, 95395, 13402, 65666, 3436, 24973, 39853, 10089, 28569, 33664, 55011, 89708, 63198, 27621, 40630, 43484, 34465, 11336, 42982, 57861, 46354, 77310, 36781, 62702, 99705, 93078, 50247, 68036, 70024, 27240, 90931, 35380, 75291, 32835, 34434, 75453, 55228, 79765, 28142, 54959, 7148, 72806, 60390, 71646, 60651, 66950, 96680, 88808, 67513, 53689, 28976, 46643, 54085, 59222, 41109, 90677, 42888, 72101, 39604, 82146, 52191, 38324, 13596, 30048, 2976, 90832, 8791, 5908, 40390, 74304, 41995, 36256, 88112, 84726, 13416, 62747, 36911, 3...
scala> q ++= nums
res5: q.type = PriorityQueue(99995, 99996, 99997, 99998, 99999)
///////////////// SAMPLE USAGE ////////////////////
scala>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment