Bounded priority queue in Scala

gistfile1.scala

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 ([email protected])
 */
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>

Alas, this no longer works: Illegal inheritance from sealed class PriorityQueue

just one more reason to stick with scala 2.11.x

https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/util/BoundedPriorityQueue.scala