How to add unzipWithIndex to all Scala collections where it makes sense

Suppose I have a list of strings and I use zipWithIndex to transform it to a list of tuples:

List("a", "b", "c").zipWithIndex
res1: List[(java.lang.String, Int)] = List((a,0), (b,1), (c,2))

I'd like to write an unzipWithIndex method that performs the reverse transformation, i.e. a method which, when applied to a list of tuples whose second elements are a permutation of the first N integers, returns the firs开发者_高级运维t elements in unpermuted order:

List(("c",2), ("a",0), ("b",1)).unzipWithIndex
res2: List[java.lang.String] = List(a, b, c)

The method should work on any suitable collection of 2-tuples whose second elements are of type Int, preferably using the Pimp My Library pattern. How would I go about this with Scala 2.8 collections?

object Test {
  import collection.generic.CanBuildFrom

  class Unzip[T, CC[X] <: Traversable[X]]
    (coll: CC[(T, Int)])
    (implicit bf: CanBuildFrom[CC[(T, Int)], T, CC[T]]) {
    def unzipWithIndex: CC[T] = bf(coll) ++= (coll.toSeq sortBy (_._2) map (_._1)) result
  }

  implicit def installUnzip[T, CC[X] <: Traversable[X]]
    (coll: CC[(T, Int)])
    (implicit bf: CanBuildFrom[CC[(T, Int)], T, CC[T]]) = new Unzip[T, CC](coll)

  def main(args: Array[String]): Unit = {
    val x1 = util.Random.shuffle("abcdefgh".zipWithIndex)

    println("x1 shuffled = " + x1)
    println("x1.unzipWithIndex = " + x1.unzipWithIndex)

    val x2 = (1 to 10).toSet.zipWithIndex

    println("x2 = " + x2)
    println("x2.unzipWithIndex = " + x2.unzipWithIndex)
  }
}

% scala Test
x1 shuffled = Vector((f,5), (g,6), (c,2), (d,3), (e,4), (a,0), (h,7), (b,1))
x1.unzipWithIndex = Vector(a, b, c, d, e, f, g, h)
x2 = Set((8,8), (2,5), (3,7), (5,0), (9,4), (4,9), (6,3), (10,1), (7,6), (1,2))
x2.unzipWithIndex = Set(5, 10, 1, 6, 9, 2, 7, 3, 8, 4)

Not an exact match, but you could draw some ideas from the unzip implementation of HList (an arbitrary-length tuple2, coming from Haskell Strongly typed heterogeneous collections).

See Type-Level Programming in Scala, Part 6d: HList Zip/Unzip from Rúnar Óli and Mark Harrah .

we’ll define an unzip type class that accepts an HList of tuples and separates it into two HLists by components:

trait Unzip[H <: HList, R1 <: HList, R2 <: HList] {
   def unzip(h: H): (R1, R2)
}

Unzipping HNil produces HNil.

implicit def unzipNil =
   new Unzip[HNil, HNil, HNil] {
      def unzip(h: HNil) = (HNil, HNil)
   }

For HCons, we unzip the tail, separate the head components, and prepend the respective head component to each tail component.

   implicit def unzipCons[H1, H2, T <: HList, TR1 <: HList, TR2 <: HList]
      (implicit unzipTail: Unzip[T, TR1, TR2]) =

      new Unzip[(H1,H2) :: T, H1 :: TR1, H2 :: TR2]  {
         def unzip(h: (H1,H2) :: T) = {
            val (t1, t2) = unzipTail.unzip(h.tail)
            (HCons(h.head._1, t1), HCons(h.head._2, t2))
         }
      }

   def unzip[H <: HList, R1 <: HList, R2 <: HList](h: H)(implicit un: Unzip[H, R1, R2]): (R1, R2) =
      un unzip h
}

Again, we just need to hook this into our HListOps type class.

Building on the example from above,

// unzip the zipped HLists
val (cc1, cc2) = cc.unzip

val (ca, cb) = cc1.unzip

Here is my take:

implicit def toUnzippable[E](l:List[(E,Int)]) = new AnyRef {
  def unzipWithIndex:List[E] = l.sortBy(_._2).map(_._1)
}

This does what your example requires. However, there are two problems with it:

• Does not handle missing indices (e.g. List(("a",0), ("b",2)).unzipWithIndex == List("a","b"))
• Only works with lists.

You can actually make it to work with all sequences that have builders, so they return the same type of sequence which you put in. But that would require some more work on it.