获取 Scala Iterable 的前 n 个元素的最简单方法

Question

是否有一个简单有效的解决方案来确定 Scala Iterable 的前 n 个元素？我的意思是类似的东西

iter.toList.sortBy(_.myAttr).take(2)

，但当只有前两个元素感兴趣时，不必对所有元素进行排序。理想情况下，我正在寻找类似的内容

iter.top(2, _.myAttr)

，另请参阅：Solution for the top element using an Ordering: In Scala, how to use Ordering[T] with List .min 或 List.max 并保持代码可读

更新：

谢谢大家的解决方案。最后，我采用了用户未知的原始解决方案，并采用它来使用Iterable和pimp-my-library模式：

implicit def iterExt[A](iter: Iterable[A]) = new {
  def top[B](n: Int, f: A => B)(implicit ord: Ordering[B]): List[A] = {
    def updateSofar (sofar: List [A], el: A): List [A] = {
      //println (el + " - " + sofar)

      if (ord.compare(f(el), f(sofar.head)) > 0)
        (el :: sofar.tail).sortBy (f)
      else sofar
    }

    val (sofar, rest) = iter.splitAt(n)
    (sofar.toList.sortBy (f) /: rest) (updateSofar (_, _)).reverse
  }
}

case class A(s: String, i: Int)
val li = List (4, 3, 6, 7, 1, 2, 9, 5).map(i => A(i.toString(), i))
println(li.top(3, _.i))

Answer 1

我的解决方案（绑定到 Int，但应该很容易更改为 Ordered（请几分钟）：

def top (n: Int, li: List [Int]) : List[Int] = {

  def updateSofar (sofar: List [Int], el: Int) : List [Int] = {
    // println (el + " - " + sofar)
    if (el < sofar.head) 
      (el :: sofar.tail).sortWith (_ > _) 
    else sofar
  }

  /* better readable:
    val sofar = li.take (n).sortWith (_ > _)
    val rest = li.drop (n)
    (sofar /: rest) (updateSofar (_, _)) */    
  (li.take (n). sortWith (_ > _) /: li.drop (n)) (updateSofar (_, _)) 
}

用法：

val li = List (4, 3, 6, 7, 1, 2, 9, 5)    
top (2, li)

• 对于上面的列表，将前 2 (4, 3) 作为起始 TopTen (TopTwo)。
• 对它们进行排序，使得第一个元素是较大的元素（如果有），
• 重复迭代列表的其余部分（li.drop(n)），并将当前元素与最小值列表中的最大值进行比较；如果需要，请替换，然后再次使用。
• 改进：
  • 丢弃 Int，并使用有序。
  • 丢弃 (_ > _) 并使用用户排序来允许 BottomTen。 （较难：选择中间的 10 个 :) ）
  • 扔掉 List，改用 Iterable

更新（抽象）：

def extremeN [T](n: Int, li: List [T])
  (comp1: ((T, T) => Boolean), comp2: ((T, T) => Boolean)):
     List[T] = {

  def updateSofar (sofar: List [T], el: T) : List [T] =
    if (comp1 (el, sofar.head)) 
      (el :: sofar.tail).sortWith (comp2 (_, _)) 
    else sofar

  (li.take (n) .sortWith (comp2 (_, _)) /: li.drop (n)) (updateSofar (_, _)) 
}

/*  still bound to Int:  
def top (n: Int, li: List [Int]) : List[Int] = {
  extremeN (n, li) ((_ < _), (_ > _))
}
def bottom (n: Int, li: List [Int]) : List[Int] = {
  extremeN (n, li) ((_ > _), (_ < _))
}
*/

def top [T] (n: Int, li: List [T]) 
  (implicit ord: Ordering[T]): Iterable[T] = {
  extremeN (n, li) (ord.lt (_, _), ord.gt (_, _))
}
def bottom [T] (n: Int, li: List [T])
  (implicit ord: Ordering[T]): Iterable[T] = {
  extremeN (n, li) (ord.gt (_, _), ord.lt (_, _))
}

top (3, li)
bottom (3, li)
val sl = List ("Haus", "Garten", "Boot", "Sumpf", "X", "y", "xkcd", "x11")
bottom (2, sl)

用Iterable替换List似乎有点难。

正如 Daniel C. Sobral 在评论中指出的那样，topN 中的高 n 可能会导致大量排序工作，因此进行手动插入排序而不是重复对整个列表进行排序可能会很有用top-n 元素的个数：

def extremeN [T](n: Int, li: List [T])
  (comp1: ((T, T) => Boolean), comp2: ((T, T) => Boolean)):
     List[T] = {

  def sortedIns (el: T, list: List[T]): List[T] = 
    if (list.isEmpty) List (el) else 
    if (comp2 (el, list.head)) el :: list else 
      list.head :: sortedIns (el, list.tail)

  def updateSofar (sofar: List [T], el: T) : List [T] =
    if (comp1 (el, sofar.head)) 
      sortedIns (el, sofar.tail)
    else sofar

  (li.take (n) .sortWith (comp2 (_, _)) /: li.drop (n)) (updateSofar (_, _)) 
}

top/bottom 方法和用法如上。对于顶部/底部元素的小组，很少调用排序，一开始会调用几次，然后随着时间的推移越来越少。例如，top (10) 为 10 000 时 70 次，top (10) 为 100 000 时 90 次。

Answer 2

这是另一个简单且性能相当好的解决方案。

def pickTopN[T](k: Int, iterable: Iterable[T])(implicit ord: Ordering[T]): Seq[T] = {
  val q = collection.mutable.PriorityQueue[T](iterable.toSeq:_*)
  val end = Math.min(k, q.size)
  (1 to end).map(_ => q.dequeue())
}

大 O 是 O(n + k log n)，其中 k <= n。因此，对于较小的 k 和最坏的 n log n 来说，性能是线性的。

该解决方案还可以优化为内存复杂度为 O(k)，性能优化为 O(n log k)。这个想法是使用 MinHeap 始终只跟踪前 k 个项目。这是解决方案。

def pickTopN[A, B](n: Int, iterable: Iterable[A], f: A => B)(implicit ord: Ordering[B]): Seq[A] = {
  val seq = iterable.toSeq
  val q = collection.mutable.PriorityQueue[A](seq.take(n):_*)(ord.on(f).reverse) // initialize with first n

  // invariant: keep the top k scanned so far
  seq.drop(n).foreach(v => {
    q += v
    q.dequeue()
  })

  q.dequeueAll.reverse
}

Answer 3

还有另一个版本：

val big = (1 to 100000)

def maxes[A](n:Int)(l:Traversable[A])(implicit o:Ordering[A]) =
    l.foldLeft(collection.immutable.SortedSet.empty[A]) { (xs,y) =>
      if (xs.size < n) xs + y
      else {
        import o._
        val first = xs.firstKey
        if (first < y) xs - first + y
        else xs
      }
    }

println(maxes(4)(big))
println(maxes(2)(List("a","ab","c","z")))

使用 Set 强制列表具有唯一值：

def maxes2[A](n:Int)(l:Traversable[A])(implicit o:Ordering[A]) =
    l.foldLeft(List.empty[A]) { (xs,y) =>
      import o._
      if (xs.size < n) (y::xs).sort(lt _)
      else {
        val first = xs.head
        if (first < y) (y::(xs - first)).sort(lt _)
        else xs
      }
    }

Answer 4

您无需对整个集合进行排序即可确定前 N 个元素。但是，我不认为此功能是由原始库提供的，因此您必须自己开发，可能使用 pimp-my-library 模式。

例如，您可以按如下方式获取集合的第 n 个元素：（

  class Pimp[A, Repr <% TraversableLike[A, Repr]](self : Repr) {

    def nth(n : Int)(implicit ord : Ordering[A]) : A = {
      val trav : TraversableLike[A, Repr] = self
      var ltp : List[A] = Nil
      var etp : List[A] = Nil
      var mtp : List[A] = Nil
      trav.headOption match {
        case None      => error("Cannot get " + n + " element of empty collection")
        case Some(piv) =>
          trav.foreach { a =>
            val cf = ord.compare(piv, a)
            if (cf == 0) etp ::= a
            else if (cf > 0) ltp ::= a
            else mtp ::= a
          }
          if (n < ltp.length)
            new Pimp[A, List[A]](ltp.reverse).nth(n)(ord)
          else if (n < (ltp.length + etp.length))
            piv
          else
            new Pimp[A, List[A]](mtp.reverse).nth(n - ltp.length - etp.length)(ord)
      }
    }
  }

这不是很有用；抱歉）

获取前 n 个元素就很简单了：

def topN(n : Int)(implicit ord : Ordering[A], bf : CanBuildFrom[Repr, A, Repr]) ={
  val b = bf()
  val elem = new Pimp[A, Repr](self).nth(n)(ord)
  import util.control.Breaks._
  breakable {
    var soFar = 0
    self.foreach { tt =>
      if (ord.compare(tt, elem) < 0) {
         b += tt
         soFar += 1
      }
    }
    assert (soFar <= n)
    if (soFar < n) {
      self.foreach { tt =>
        if (ord.compare(tt, elem) == 0) {
          b += tt
          soFar += 1
        }
        if (soFar == n) break
      }
    }

  }
  b.result()
}

不幸的是，我在获取此皮条客时遇到了麻烦通过这个隐式发现：

implicit def t2n[A, Repr <% TraversableLike[A, Repr]](t : Repr) : Pimp[A, Repr] 
  = new Pimp[A, Repr](t)

我明白了：

scala> List(4, 3, 6, 7, 1, 2, 8, 5).topN(4)
<console>:9: error: could not find implicit value for evidence parameter of type (List[Int]) => scala.collection.TraversableLike[A,List[Int]]
   List(4, 3, 6, 7, 1, 2, 8, 5).topN(4)
       ^

但是，代码实际上工作正常：

scala> new Pimp(List(4, 3, 6, 7, 1, 2, 8, 5)).topN(4)
res3: List[Int] = List(3, 1, 2, 4)

并且

scala> new Pimp("ioanusdhpisjdmpsdsvfgewqw").topN(6)
res2: java.lang.String = adddfe

Answer 5

如果目标是不对整个列表进行排序，那么您可以执行类似的操作（当然可以稍微优化一下，以便当数字显然不应该存在时我们不会更改列表）：

List(1,6,3,7,3,2).foldLeft(List[Int]()){(l, n) => (n :: l).sorted.take(2)}

Answer 6

我最近在 Rank Apache Jackrabbit 类（不过在 Java 中）。请参阅 take 方法了解其要点。基本思想是快速排序，但一旦找到前 n 个元素就提前终止。

Answer 7

这是渐近 O(n) 解。

def top[T](data: List[T], n: Int)(implicit ord: Ordering[T]): List[T] = {
    require( n < data.size)

    def partition_inner(shuffledData: List[T], pivot: T): List[T] = 
      shuffledData.partition( e => ord.compare(e, pivot) > 0 ) match {
          case (left, right) if left.size == n => left
          case (left, x :: rest) if left.size < n => 
            partition_inner(util.Random.shuffle(data), x)
          case (left @ y :: rest, right) if left.size > n => 
            partition_inner(util.Random.shuffle(data), y)
      }

     val shuffled = util.Random.shuffle(data)
     partition_inner(shuffled, shuffled.head)
}

scala> top(List.range(1,10000000), 5)

由于递归，此解决方案将比上面的一些非线性解决方案花费更长的时间，并且可能导致java.lang.OutOfMemoryError：超出GC开销限制。
但恕我直言，它的可读性和功能性风格稍好一些。仅用于工作面试；）。

更重要的是，该解决方案可以轻松并行化。

def top[T](data: List[T], n: Int)(implicit ord: Ordering[T]): List[T] = {
    require( n < data.size)

    @tailrec
    def partition_inner(shuffledData: List[T], pivot: T): List[T] = 
      shuffledData.par.partition( e => ord.compare(e, pivot) > 0 ) match {
          case (left, right) if left.size == n => left.toList
          case (left, right) if left.size < n => 
            partition_inner(util.Random.shuffle(data), right.head)
          case (left, right) if left.size > n => 
            partition_inner(util.Random.shuffle(data), left.head)
      }

     val shuffled = util.Random.shuffle(data)
     partition_inner(shuffled, shuffled.head)
}

Answer 8

对于 n 的较小值和较大的列表，可以通过选取最大元素 n 次来获取前 n 元素：

def top[T](n:Int, iter:Iterable[T])(implicit ord: Ordering[T]): Iterable[T] = {
  def partitionMax(acc: Iterable[T], it: Iterable[T]): Iterable[T]  = {
    val max = it.max(ord)
    val (nextElems, rest) = it.partition(ord.gteq(_, max))
    val maxElems = acc ++ nextElems
    if (maxElems.size >= n || rest.isEmpty) maxElems.take(n)
    else partitionMax(maxElems, rest)
  }
  if (iter.isEmpty) iter.take(0)
  else partitionMax(iter.take(0), iter)
}

这并不对整个列表进行排序并采用Ordering。我相信我在 partitionMax 中调用的每个方法都是 O（列表大小），并且我只希望最多调用它 n 次，因此小 n 的整体效率 将与迭代器的大小成正比。

scala> top(5, List.range(1,1000000))
res13: Iterable[Int] = List(999999, 999998, 999997, 999996, 999995)

scala> top(5, List.range(1,1000000))(Ordering[Int].on(- _))
res14: Iterable[Int] = List(1, 2, 3, 4, 5)

您还可以在 n 接近可迭代的大小时添加一个分支，并切换到 iter.toList.sortBy(_.myAttr).take(n)。

它不会返回提供的集合类型，但您可以查看 如果有要求，如何将丰富我的库模式应用于 Scala 集合？。

Answer 9

使用 PriorityQueue 的优化解决方案，时间复杂度为 O(nlogk)。在更新中给出的方法中，您每次都会对 sofar 列表进行排序，这是不需要的，并且在其下方使用 PriorityQueue 进行优化。

import scala.language.implicitConversions
import scala.language.reflectiveCalls
import collection.mutable.PriorityQueue
implicit def iterExt[A](iter: Iterable[A]) = new {
    def top[B](n: Int, f: A => B)(implicit ord: Ordering[B]) : List[A] = {
        def updateSofar (sofar: PriorityQueue[A], el: A): PriorityQueue[A] = {
            if (ord.compare(f(el), f(sofar.head)) < 0){
                sofar.dequeue
                sofar.enqueue(el)
            }
            sofar
        }

        val (sofar, rest) = iter.splitAt(n)
        (PriorityQueue(sofar.toSeq:_*)( Ordering.by( (x :A) => f(x) ) ) /: rest) (updateSofar (_, _)).dequeueAll.toList.reverse
    }
}

case class A(s: String, i: Int)
val li = List (4, 3, 6, 7, 1, 2, 9, 5).map(i => A(i.toString(), i))
println(li.top(3, -_.i))