impl top K items

This commit is contained in:
Ivan Li 2019-01-10 13:35:55 +08:00
parent 66857c1fbc
commit ab9f58704c
2 changed files with 55 additions and 0 deletions

View File

@ -0,0 +1,30 @@
package ch29_heap_solutions
import scala.collection.mutable
/**
* keep the top k items in the the class
*/
class TopKItemsKeeper(val itemsToKeepCount: Int) {
//have a smallest value top heap
val queue = new mutable.PriorityQueue[Int]()(scala.math.Ordering.Int.reverse)
def put(item: Int): Unit = {
if (queue.length < itemsToKeepCount) {
queue += item
return
}
//queue already have the k items
if (item.compareTo(queue.head) > 0) {
queue.dequeue()
queue += item
}
}
def get(): List[Int] = {
queue.clone().dequeueAll
}
}

View File

@ -0,0 +1,25 @@
package ch29_heap_solutions
import org.scalatest.{FlatSpec, Matchers}
import scala.util.Random
class TopKItemsKeeperTest extends FlatSpec with Matchers {
behavior of "TopKItemsKeeperTest"
it should "put and get top K from the keeper" in {
val length = 50
val k = 5
val topKItemsKeeper = new TopKItemsKeeper(k)
val nums = new Array[Int](length)
for (i <- Range(0, length)) {
nums(i) = Random.nextInt
}
nums.foreach(topKItemsKeeper.put)
val ordering = scala.math.Ordering.Int.reverse
topKItemsKeeper.get().toArray.sorted(ordering) should equal(nums.sorted(ordering).slice(0, k))
}
}