Top N Sorter in Scala

Scala: トップNソートの実装

 

比較可能な大きな集合 $M$ に対して上位 $N$ 件のソート結果を取得したい。

この局面では一般にヒープソートが有効なので、PriorityQueue を使ってみる。

 

import scala.collection.mutable.PriorityQueue

class TopNSorter[A](n: Int)(ord: Ordering[A]) {
  private[this] val data = PriorityQueue()(ord.reverse)

  def insert(item: A): Unit = {
    data += item
    if (data.size > n) data.dequeue
  }

  def result(): Seq[A] = data.dequeueAll.reverse
}

insert は $O(\log N)$ になるはず。(要素の追加は $O(\log N)$, サイズ取得と先頭要素へのアクセスは$O(1)$)
$M$ 件挿入すれば $O(M\log N)$。
result は $O(N)$ だが、$N$ が小さい前提なので問題ない。

scala> val x = new TopNSorter[Int](100)(Ordering.Int)
x: TopNSorter[Int] = TopNSorter@520f1862

scala> (1 to 10000000) foreach x.insert

scala> x.result
res1: Seq[Int] = Vector(10000000, 9999999, 9999998, 9999997, 9999996, 9999995, 9999994, 9999993, 9999992, 9999991, 9999990, 9999989, 9999988, 9999987, 9999986, 9999985, 9999984, 9999983, 9999982, 9999981, 9999980, 9999979, 9999978, 9999977, 9999976, 9999975, 9999974, 9999973, 9999972, 9999971, 9999970, 9999969, 9999968, 9999967, 9999966, 9999965, 9999964, 9999963, 9999962, 9999961, 9999960, 9999959, 9999958, 9999957, 9999956, 9999955, 9999954, 9999953, 9999952, 9999951, 9999950, 9999949, 9999948, 9999947, 9999946, 9999945, 9999944, 9999943, 9999942, 9999941, 9999940, 9999939, 9999938, 9999937, 9999936, 9999935, 9999934, 9999933, 9999932, 9999931, 9999930, 9999929, 9999928, 9999927, 9999926, 9999925, 9999924, 9999923, 9999922, 9999921, 9999920, 9999919, 9999918, 9999917, 9999916, 9999915...

単に

scala> (1 to 10000000).sorted.reverse.take(100)

とするよりはずっと速い。

もう少しスマートに書けそうなものだけど。。。