How to create a generic groupByIndex function? - scala

Question

I would like to have groupByIndex function that groups values based on their index (and not the value).
A concrete method definition for Vector[A] could look like the following:
def groupByIndex[A, K](vector: Vector[A], f: Int => K): immutable.Map[K, Vector[(A, Int)]] = {
vector.zipWithIndex.groupBy { case (elem, index) => f(index) }
}
Testing this function in the REPL gives indeed the correct result:
scala> val vec = Vector.tabulate(4)(i => s"string ${i+1}")
vec: scala.collection.immutable.Vector[String] = Vector(string 1, string 2, string 3, string 4)
scala> groupByIndex(vec, i => i%2)
res2: scala.collection.immutable.Map[Int,Vector[(String, Int)]] = Map(1 -> Vector((string 2,1), (string 4,3)), 0 -> Vector((string 1,0), (string 3,2)))
Now, I would like to apply the "enrich-my-library" pattern to give this method to all the classes that should support it, i.e. classes that implement zipWithIndex and groupBy. Those two methods are defined in GenIterableLike (zipWithIndex) and GenTraversableLike/TraversableLike (groupBy).
With all this in mind, I tried to mimic the method definitions of zipWithIndex (this is the problematic) and groupBy to build my own groupByIndex:
implicit class GenIterableLikeOps[A, Repr](val iterable: GenIterableLike[A, Repr] with TraversableLike[A, Repr]) extends AnyVal {
def groupByIndex[K, A1 >: A, That <: TraversableLike[(A1, Int), OtherRepr], OtherRepr](f: Int => K)(implicit bf: CanBuildFrom[Repr, (A1, Int), That]): immutable.Map[K, OtherRepr] = {
val zipped = iterable.zipWithIndex
zipped.groupBy{ case (elem, index) => f(index) }
}
}
First, this seems way too complicated to me - is there a way to simplify this? For example, can we somehow drop the second OtherRepr? (I was not able to.)
Second, I am not able to call this function without explicitly specifying the generic parameters. Using the example from above I get the following error:
scala> vec.groupByIndex(i => i%2)
<console>:21: error: Cannot construct a collection of type scala.collection.TraversableLike[(String, Int),Nothing] with elements of type (String, Int) based on a collection of type scala.collection.immutable.Vector[String].
vec.groupByIndex(i => i%2)
^
scala> vec.groupByIndex[Int, String, Vector[(String, Int)], Vector[(String, Int)]](i => i%2)
res4: scala.collection.immutable.Map[Int,Vector[(String, Int)]] = Map(1 -> Vector((string 2,1), (string 4,3)), 0 -> Vector((string 1,0), (string 3,2)))
How do I a) simplify this method and b) make it work without having to specify the generic parameters?

Answer 1

You can substitute the OtherThat type parameter by That. That way you get rid of OtherThat and solve the problem of having to specify the generic type parameters. The compiler is then able to resolve That by looking at the implicit value for CanBuildFrom[Repr, (A1, Int), That].
implicit class GenIterableLikeOps[A, Repr]
(val iterable: GenIterableLike[A, Repr] with TraversableLike[A, Repr])
extends AnyVal {
def groupByIndex
[K, A1 >: A, That <: TraversableLike[(A1, Int), That]]
(f: Int => K)(implicit bf: CanBuildFrom[Repr, (A1, Int), That])
: Map[K, That] = {
val zipped = iterable.zipWithIndex
zipped.groupBy{ case (elem, index) => f(index) }
}
}

Answer 2

This isn't as good as the other answer, but if you don't care about what you're building, one way to simplify and avoid building the zipped collection:
implicit class gbi[A](val as: Traversable[A]) extends AnyVal {
def groupByIndex[K](f: Int => K) = (as, (0 until Int.MaxValue)).zipped.groupBy { case (x, i) => f(i) }
}
The range is a benign way to avoid taking the size of the traversable.