Scala で zipWith メソッドを使う

scala の collection はなぜか zipWith メソッドを持っていない。
zipWith を実装しなくても

(xs zip ys) map { case (x, y) => op(x, y) }
(xs, ys).zipped map op

などで同じことができるけれど、やっぱり zipWith が使いたい。
じゃあ作ってみよう、というのがこの記事のテーマです。

とりあえず作ってみる

implicit class を使うとメソッドの追加みたいなことができる。
こんな感じでどうだろうか。

implicit class RichIterable[A](xs: Iterable[A]) {
  def zipWith[B, C](ys: Iterable[B])(op: (A, B) => C): Iterable[C] = {
    (xs zip ys) map { case (x, y) => op(x, y) }
  }
}

repl で試してみる。

$ scala
Welcome to Scala version 2.11.6 (Java HotSpot(TM) 64-Bit Server VM, Java 1.8.0_45).
Type in expressions to have them evaluated.
Type :help for more information.

scala> implicit class RichIterable[A](xs: Iterable[A]) {
     |   def zipWith[B, C](ys: Iterable[B])(op: (A, B) => C): Iterable[C] = {
     |     (xs zip ys) map { case (x, y) => op(x, y) }
     |   }
     | }
defined class RichIterable

scala> val a = Set(1, 2, 3, 4)
a: scala.collection.immutable.Set[Int] = Set(1, 2, 3, 4)

scala> val b = List(0.1, 0.2, 0.3)
b: List[Double] = List(0.1, 0.2, 0.3)

scala> a.zipWith(b)(_+_)
res0: Iterable[Double] = Set(1.1, 2.2, 3.3)

出来た！ と思ったら型が Iterable[Double] になってしまっている。
もちろん zip, zipped, map を用いた場合は、ちゃんと型が Set[Double] となる。

scala> (a zip b) map { case (x, y) => x + y }
res1: scala.collection.immutable.Set[Double] = Set(1.1, 2.2, 3.3)

scala> (a, b).zipped map (_+_)
res2: scala.collection.immutable.Set[Double] = Set(1.1, 2.2, 3.3)

型が Set[Double] になってくれないと、

scala> val c: Set[Double] = a.zipWith(b)(_+_)
<console>:10: error: type mismatch;
 found   : Iterable[Double]
 required: Set[Double]
       val c: Set[Double] = a.zipWith(b)(_+_)

こんなことになってしまう。
ということで、

map と zip を作る

調べてみると、map と zip は

TraversableLike[A, Repr] に対して
map[B, That](f: (A) ⇒ B)(implicit bf: CanBuildFrom[Repr, B, That]): That

IterableLike[A, Repr] に対して
zip[A1 >: A, B, That](that: GenIterable[B])(implicit bf: CanBuildFrom[Repr, (A1, B), That]): That

となっている。
ん、なんで zip は A1 が必要なんだろう？
よくわからないけれど、とりあえず実装してみる。

mapzip.scala

import scala.collection.generic.CanBuildFrom
import scala.collection.{TraversableLike, IterableLike, GenIterable}

implicit class RichTraversableLike[A, Repr](xs: TraversableLike[A, Repr]) {
  def myMap[B, That](f: A => B)(implicit bf: CanBuildFrom[Repr, B, That]): That = xs map f
}

implicit class RichIterableLike[A, Repr](xs: IterableLike[A, Repr]) {
  def myZip1[A1 >: A, B, That](that: GenIterable[B])(implicit bf: CanBuildFrom[Repr, (A1, B), That]): That = xs zip that
  def myZip2[B, That](that: GenIterable[B])(implicit bf: CanBuildFrom[Repr, (A, B), That]): That = xs zip that
}

repl で試してみる。

scala> :load mapzip.scala
Loading mapzip.scala...
import scala.collection.generic.CanBuildFrom
import scala.collection.{TraversableLike, IterableLike, GenIterable}
defined class RichTraversableLike
defined class RichIterableLike

scala> a.myMap(_+1)
res3: scala.collection.immutable.Set[Int] = Set(2, 3, 4, 5)

scala> a.myZip1(b)
res4: scala.collection.immutable.Set[(Int, Double)] = Set((1,0.1), (2,0.2), (3,0.3))

scala> a.myZip2(b)
res5: scala.collection.immutable.Set[(Int, Double)] = Set((1,0.1), (2,0.2), (3,0.3))

よし、良さそうだ。
myZip2 でも問題ないように思える。

zipWith を作る。

作った zip と map を組み合わせればできるはず。
こんな感じになった。

zipWith.scala

import scala.collection.generic.CanBuildFrom
import scala.collection.{TraversableLike, IterableLike, GenIterable}

implicit class RichIterableLike[A, Repr](xs: IterableLike[A, Repr]) {
  def zipWith[B, C, Repr2, That](ys: GenIterable[B])(op: (A, B) => C)(implicit bf1: CanBuildFrom[Repr, (A, B), TraversableLike[(A, B), Repr2]], bf2: CanBuildFrom[Repr2, C, That]): That = {
    (xs zip ys) map { case (x, y) => op(x, y) }
  }
}

repl で試してみる。
前の zipWith が残っているのでいったん repl を終了させてからもう一度起動する。

$ scala
Welcome to Scala version 2.11.6 (Java HotSpot(TM) 64-Bit Server VM, Java 1.8.0_45).
Type in expressions to have them evaluated.
Type :help for more information.

scala> :load zipWith.scala
Loading zipWith.scala...
import scala.collection.generic.CanBuildFrom
import scala.collection.{TraversableLike, IterableLike, GenIterable}
defined class RichIterableLike

scala> val a = Set(1, 2, 3, 4)
a: scala.collection.immutable.Set[Int] = Set(1, 2, 3, 4)

scala> val b = List(0.1, 0.2, 0.3)
b: List[Double] = List(0.1, 0.2, 0.3)

scala> a.zipWith(b)(_+_)
res0: scala.collection.immutable.Set[Double] = Set(1.1, 2.2, 3.3)

今度は大丈夫そうだ。
ちゃんと型が Set[Double] になっている。
これで zipWith が使える！