How can I extend Scala collections with an argmax method?

Question

I would like to add to all collections where it makes sense, an argMax method. How to do it? Use implicits?

Answer 1

You can add functions to an existing API in Scala by using the Pimp my Library pattern. You do this by defining an implicit conversion function. For example, I have a class Vector3 to represent 3D vectors:

class Vector3 (val x: Float, val y: Float, val z: Float)

Suppose I want to be able to scale a vector by writing something like: 2.5f * v. I can't directly add a * method to class Float ofcourse, but I can supply an implicit conversion function like this:

implicit def scaleVector3WithFloat(f: Float) = new {
    def *(v: Vector3) = new Vector3(f * v.x, f * v.y, f * v.z)
}

Note that this returns an object of a structural type (the new { ... } construct) that contains the * method.

I haven't tested it, but I guess you could do something like this:

implicit def argMaxImplicit[A](t: Traversable[A]) = new {
    def argMax() = ...
}

Answer 2

Yes, the usual way would be to use the 'pimp my library' pattern to decorate your collection. For example (N.B. just as illustration, not meant to be a correct or working example):

trait PimpedList[A] {
   val l: List[A]

  //example argMax, not meant to be correct
  def argMax[T <% Ordered[T]](f:T => T) = {error("your definition here")}
}

implicit def toPimpedList[A](xs: List[A]) = new PimpedList[A] { 
  val l = xs
}

scala> def f(i:Int):Int = 10
f: (i: Int) Int

scala> val l = List(1,2,3)
l: List[Int] = List(1, 2, 3)

scala> l.argMax(f)
java.lang.RuntimeException: your definition here
    at scala.Predef$.error(Predef.scala:60)
    at PimpedList$class.argMax(:12)
        //etc etc...

Answer 3

On Scala 2.8, this works:

val list = List(1, 2, 3)
def f(x: Int) = -x
val argMax = list max (Ordering by f)

As pointed by mkneissl, this does not return the set of maximum points. Here's an alternate implementation that does, and tries to reduce the number of calls to f. If calls to f don't matter that much, see mkneissl's answer. Also, note that his answer is curried, which provides superior type inference.

def argMax[A, B: Ordering](input: Iterable[A], f: A => B) = {
  val fList = input map f
  val maxFList = fList.max
  input.view zip fList filter (_._2 == maxFList) map (_._1) toSet
}

scala> argMax(-2 to 2, (x: Int) => x * x)
res15: scala.collection.immutable.Set[Int] = Set(-2, 2)

Answer 4

The argmax function (as I understand it from Wikipedia)

def argMax[A,B](c: Traversable[A])(f: A=>B)(implicit o: Ordering[B]): Traversable[A] = {
  val max = (c map f).max(o)
  c filter { f(_) == max }
}

If you really want, you can pimp it onto the collections

implicit def enhanceWithArgMax[A](c: Traversable[A]) = new {
  def argMax[B](f: A=>B)(implicit o: Ordering[B]): Traversable[A] = ArgMax.argMax(c)(f)(o)
}

and use it like this

val l = -2 to 2
assert (argMax(l)(x => x*x) == List(-2,2))
assert (l.argMax(x => x*x) == List(-2,2))

(Scala 2.8)

Answer 5

Here's a way of doing so with the implicit builder pattern. It has the advantage over the previous solutions that it works with any Traversable, and returns a similar Traversable. Sadly, it's pretty imperative. If anyone wants to, it could probably be turned into a fairly ugly fold instead.

object RichTraversable {
  implicit def traversable2RichTraversable[A](t: Traversable[A]) = new RichTraversable[A](t)
}

class RichTraversable[A](t: Traversable[A]) { 
  def argMax[That, C](g: A => C)(implicit bf : scala.collection.generic.CanBuildFrom[Traversable[A], A, That], ord:Ordering[C]): That = {
    var minimum:C = null.asInstanceOf[C]
    val repr = t.repr
    val builder = bf(repr)
    for(a<-t){
      val test: C = g(a)
      if(test == minimum || minimum == null){
        builder += a
        minimum = test
      }else if (ord.gt(test, minimum)){
        builder.clear
        builder += a
        minimum = test
      }
    }
    builder.result
  }
}

Set(-2, -1, 0, 1, 2).argmax(x=>x*x) == Set(-2, 2)
List(-2, -1, 0, 1, 2).argmax(x=>x*x) == List(-2, 2)