Updating a Map in Scala

Relevant references:

Let’s say we want to count words with the help of a map:

val m1 = Map[String, Int]()

Adding a new word to the map is easy:

val m2 = m1.updated("foo", 1)

val m2alt = m1 + Pair("foo", 1)

But how can we increase the count if the word is already in the map?

Using Only Built-in Methods

A simplistic approach would be to use two separate branches:

val m3 = if (m2.contains("foo")) m2.updated("foo", m2("foo") + 1)
         else m2.updated("foo", 1)

val m3alt = m2.updated("foo", (if (m2.contains("foo")) m2("foo") else 0) + 1)

But we can surely do better by using get that returns an Option:

val m4 = m2.updated("foo", m2.get("foo").getOrElse(0) + 1)

Another approach would be to set a default value for the map:

val m2d = m2.withDefault(_ => 0)
val m5 = m2d.updated("foo", m2d("foo") + 1)

(One drawback with using default values is that you can only have a single default value strategy per map.)

The above solutions are — as far as I know — the best possible using only built-in methods.

The updatedWith Extension Method

If we like the m4 solution (I do!), we can turn it into an extension method:

implicit class MapExtensions[K, V](val map: Map[K, V]) extends AnyVal {
  def updatedWith(key: K, default: V)(f: V => V) = {
    map.updated(key, f(map.getOrElse(key, default)))
  }
}

val m6 = m2.updatedWith("foo", 0) { _ + 1 }

Scalaz

There are already extension methods like this in the library Scalaz1.

import scalaz.Scalaz._

The first one is called insertWith and takes a function that handles “collisions”:

val m7 = m2.insertWith("foo", 1) { _ + _ }

The second is alter and it works with Options (and can also remove pairs if the function returns None):

val m8 = m2.alter("foo") { x => Some(x.getOrElse(0) + 1) }

Here, I used the same getOrElse(0) as before, but since we’re now in Scalaz land we can use the fact that integers are monoids2 and have a zero provided for us:

val m9 = m2.alter("foo") { x => Some(x.orZero + 1) }

Have I missed another simple solution? Please get in touch! :)

1 Scalaz is getting a bad rap from some parts of the Scala community for being too abstract and having funny operators. While I’m sympathetic to these complaints — abstract abstractions are not mainstream — there are certainly some goodies in there.

2 This is just a fancy way of saying that integers have an associative binary operation (+) and a zero (or identity element) that fulfill some simple laws.