Simplest way to get the top n elements of a Scala Iterable

前端 未结 9 673
佛祖请我去吃肉
佛祖请我去吃肉 2020-11-29 02:48

Is there a simple and efficient solution to determine the top n elements of a Scala Iterable? I mean something like

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


        
9条回答
  •  爱一瞬间的悲伤
    2020-11-29 03:14

    You don't need to sort the entire collection in order to determine the top N elements. However, I don't believe that this functionality is supplied by the raw library, so you would have to roll you own, possibly using the pimp-my-library pattern.

    For example, you can get the nth element of a collection as follows:

      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)
          }
        }
      }
    

    (This is not very functional; sorry)

    It's then trivial to get the top n elements:

    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()
    }
    

    Unfortunately I'm having trouble getting this pimp to be discovered via this implicit:

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

    I get this:

    scala> List(4, 3, 6, 7, 1, 2, 8, 5).topN(4)
    :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)
           ^
    

    However, the code actually works OK:

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

    And

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

提交回复
热议问题