[SOLVED] How can I extend Scala collections with an argmax method?

How can I extend Scala collections with an argmax method?

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

Solution

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)