Nice way to add number to element in Scala map if key exists or insert new element it not

Question

I know about couple of similar questions. They don't help me - code does not work if there is no existing key.

I need just some nice approach to append Map with value adding it to existing key (if it does exist) or putting as NEW key (if map does not contain appropriate key).

Following code works but I don't like it:

val a = collection.mutable.Map(("k1" -> 1), ("k2" -> 5))
val key = "k1"

val elem = a.get(key)
if (elem == None) {
    a += ("k5" -> 200)
} else {
    a.update(key, elem.get + 5)
}

Any point to better one? Current Scala version is 2.10.4 and I cannot currently switch to 2.11. Mutable map is not 100% limitation but preferred.

Here is, for example, similar question but I also need to account case of non-existing key which is not accounted there. At least we should understand a.get(key) could be None or add some better approach. Good idea was |+| but I'd like to keep basic Scala 2.10.x.

Answer 1

The shortest way to do that:

a += a.get(key).map(x => key -> (x + 5)).getOrElse("k5" -> 200)

In general:

a += a.get(k).map(f).map(k -> _).getOrElse(kv)

Same if your dictionary is immutable:

m + m.get(k).map(f).map(k -> _).getOrElse(kv)

so I don't see any reason to use mutable collection here.

If you don't like all these Option.map things:

m + (if (m.contains(k)) k -> f(m(k)) else kv)

Note, that there is a whole class of possible variations:

k1 -> f(m(k1)) else k2 -> v2 //original
k1 -> f(m(k1)) else k1 -> v2
k1 -> f(m(k2)) else k2 -> v2
k1 -> f(m(k2)) else k1 -> v2
k2 -> v2 else k1 -> f(m(k1)) 
k1 -> v2 else k1 -> f(m(k1))
k2 -> v2 else k1 -> f(m(k2))
k1 -> v2 else k1 -> f(m(k2))
... //v2 may also be a function from some key's value

So, Why it's not a standard function? IMO, because all variations still can be implemented as one-liners. If you want library with all functions, that can be implemented as one-liners, you know, it's Scalaz :).

P.S. If yu also wonder why there is no "update(d) if persist" function - see @Rex Kerr 's answer here

Answer 2

You can create you own function for that purpose:

def addOrUpdate[K, V](m: collection.mutable.Map[K, V], k: K, kv: (K, V), 
                      f: V => V) {
  m.get(k) match {
    case Some(e) => m.update(k, f(e))
    case None    => m += kv
  }
}

addOrUpdate(a, "k1", "k5" -> 200, (v: Int) => v + 5)