How can I extend Scala collections with an argmax method? - scala

Question

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

Answer 1

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 2

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 3

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 4

Nice and easy ? :
val l = List(1,0,10,2)
l.zipWithIndex.maxBy(x => x._1)._2

Answer 5

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 6

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)

Answer 7

Here's a variant loosely based on #Daniel's accepted answer that also works for Sets.
def argMax[A, B: Ordering](input: GenIterable[A], f: A => B) : GenSet[A] = argMaxZip(input, f) map (_._1) toSet
def argMaxZip[A, B: Ordering](input: GenIterable[A], f: A => B): GenIterable[(A, B)] = {
if (input.isEmpty) Nil
else {
val fPairs = input map (x => (x, f(x)))
val maxF = fPairs.map(_._2).max
fPairs filter (_._2 == maxF)
}
}
One could also do a variant that produces (B, Iterable[A]), of course.

Answer 8

Based on other answers, you can pretty easily combine the strengths of each (minimal calls to f(), etc.). Here we have an implicit conversion for all Iterables (so they can just call .argmax() transparently), and a stand-alone method if for some reason that is preferred. ScalaTest tests to boot.
class Argmax[A](col: Iterable[A]) {
def argmax[B](f: A => B)(implicit ord: Ordering[B]): Iterable[A] = {
val mapped = col map f
val max = mapped max ord
(mapped zip col) filter (_._1 == max) map (_._2)
}
}
object MathOps {
implicit def addArgmax[A](col: Iterable[A]) = new Argmax(col)
def argmax[A, B](col: Iterable[A])(f: A => B)(implicit ord: Ordering[B]) = {
new Argmax(col) argmax f
}
}
class MathUtilsTests extends FunSuite {
import MathOps._
test("Can argmax with unique") {
assert((-10 to 0).argmax(_ * -1).toSet === Set(-10))
// or alternate calling syntax
assert(argmax(-10 to 0)(_ * -1).toSet === Set(-10))
}
test("Can argmax with multiple") {
assert((-10 to 10).argmax(math.pow(_, 2)).toSet === Set(-10, 10))
}
}