Scala collections: array groupBy and return array indexes for each group

Question

I have an array, something like that:

val a = Array("a", "c", "c", "z", "c", "b", "a")

and I want to get a map with keys of all different values of this array and values with a collection of relevant indexes for each such group, i.e. for a given array the answer would be:

Map(
  "a" -> Array(0, 6),
  "b" -> Array(5),
  "c" -> Array(1, 2, 4),
  "z" -> Array(3)
)

Surprisingly, it proved to be somewhat more complicated that I've anticipated. The best I've came so far with is:

a.zipWithIndex.groupBy {
  case(cnt, idx) => cnt
}.map {
  case(cnt, arr) => (cnt, arr.map {
    case(k, v) => v
  }
}

which is not either concise or easy to understand. Any better ideas?

Answer 1

Your code can be rewritten as oneliner, but it looks ugly.

as.zipWithIndex.groupBy(_._1).mapValues(_.map(_._2))

Another way is to use mutable.MultiMap

import collection.mutable.{ HashMap, MultiMap, Set }

val as = Array("a", "c", "c", "z", "c", "b", "a")
val mm = new HashMap[String, Set[Int]] with MultiMap[String, Int]

and then just add every binding

as.zipWithIndex foreach (mm.addBinding _).tupled    
//mm = Map(z -> Set(3), b -> Set(5), a -> Set(0, 6), c -> Set(1, 2, 4))

finally you can convert it mm.toMap if you want immutable version.

Answer 2

Here's a version with foldRight. I think it's reasonably clear.

val a = Array("a", "c", "c", "z", "c", "b", "a") 
a
 .zipWithIndex
 .foldRight(Map[String, List[Int]]())
            {case ((e,i), m)=> m updated (e, i::m.getOrElse(e, Nil))}
//> res0: scala.collection.immutable.Map[String,List[Int]] = Map(a -> List(0, 6)
//| , b -> List(5), c -> List(1, 2, 4), z -> List(3))