インデックス付きforeachほか

Scala イディオム

概要

先日Javaで書いたインデックス付きforeachループ
http://d.hatena.ne.jp/kya-zinc/20100107/1262856312
をScalaで書いてみた。
(なお、Scalaのfor-comprehensionはループ以外にも使える。)
Iterableのデフォルト実装は便利。

実装

ScalaのIterableには(Javaのインターフェースと異なり)abstractでない関数の実装があるので
elements(JavaのIterable#iterator()に対応)を実装するだけで
foreach, map, flatMap, filterが期待通りに動く。
リッチだ。

object IndexedForeach {
  def main(args : Array[String]) : Unit = {
    val list = List(100, 200, 300)
    val ilist = withIndex(list)
    val stream = Stream.from(101)
    val istream = withIndex(stream)

    //foreach
    for((i, n) <- ilist){
      println(n + " @ " + i + "ばんめ")
    }
    
    //map
    val list2 = ilist.map {case (i, n) => (i + n)}
    println(list2)
    
    //flatMap
    val list3 = ilist.flatMap {case (i, n) => List(i + n, i * n)}
    println(list3)
    
    //filter
    val list4 = ilist.filter {case (i, n) => i % 2 == 0}
    println(list4)
    
    //flatMap, filter, foreach
    for((i, x) <- ilist; (j, y) <- ilist; if (i + j) % 2 == 0){
      printf("(%d, %d) @ (%d, %d)ばんめ\n", x, y, i, j)
    }
    
    //foreach(Stream)
    for((i, n) <- istream){
      if(i > 5) return
      println(n + " @ " + i + "ばんめ")
    }
  }

  class IndexedIterator[A](iter : Iterator[A]) extends Iterator[(Int, A)] {
    var idx = 0
    
    override def hasNext = iter.hasNext
    
    override def next() = {
      val pair = (idx, iter.next)
      idx = idx + 1
      pair
    }
  }
  
  def withIndex[A](iter : Iterable[A]) : Iterable[(Int, A)] = 
    new Iterable[(Int, A)](){
      override def elements = new IndexedIterator[A](iter.elements)
  }
}

実行結果

100 @ 0ばんめ
200 @ 1ばんめ
300 @ 2ばんめ
ArrayBuffer(100, 201, 302)
ArrayBuffer(100, 0, 201, 200, 302, 600)
ArrayBuffer((0,100), (2,300))
(100, 100) @ (0, 0)ばんめ
(100, 300) @ (0, 2)ばんめ
(200, 200) @ (1, 1)ばんめ
(300, 100) @ (2, 0)ばんめ
(300, 300) @ (2, 2)ばんめ
101 @ 0ばんめ
102 @ 1ばんめ
103 @ 2ばんめ
104 @ 3ばんめ
105 @ 4ばんめ
106 @ 5ばんめ

メモ

  • デフォルト実装に頼るとmap, flatMap, filterの結果がArrayBufferになる
  • 予想に反してStreamでもforeachが動く
    • 無限だった場合、returnや例外でしか止められないが…
  • (A, B) => CはFunction2[A,B,C]でPair[A,B] => Cとは違う
    • for-comprehensionではパターンが使えるので(i, n) <- ...と書けるが、そうでない場合(i, n)を束縛するにはcaseが必要

改善

ArrayBufferが気に入らないので
foreach, map, flatMap, filterを適当に実装。(割と自明)
filterで迷ったが、mapを経由すれば一発だと気付く。
オブジェクトを返さないforeachをオーバーライドする意味はないかもしれない。
とりあえずArrayBufferだったところをListにすることに成功。

object IndexedForeach2 {
  def main(args : Array[String]) : Unit = {
    val list = List(100, 200, 300)
    val ilist = withIndex(list)
    val stream = Stream.from(101)
    val istream = withIndex(stream)

    //foreach
    for((i, n) <- ilist){
      println(n + " @ " + i + "ばんめ")
    }
    
    //map
    val list2 = ilist.map {case (i, n) => (i + n)}
    println(list2)
    
    //flatMap
    val list3 = ilist.flatMap {case (i, n) => List(i + n, i * n)}
    println(list3)
    
    //filter
    val list4 = ilist.filter {case (i, n) => i % 2 == 0}
    println(list4)
    
    //flatMap, filter, foreach
    for((i, x) <- ilist; (j, y) <- ilist; if (i + j) % 2 == 0){
      printf("(%d, %d) @ (%d, %d)ばんめ\n", x, y, i, j)
    }
    
    //foreach(Stream)
    for((i, n) <- istream){
      if(i > 5) return
      println(n + " @ " + i + "ばんめ")
    }
  }
  
  class IndexedIterator[A](iter : Iterator[A]) extends Iterator[(Int, A)] {
    var idx = 0
    
    override def hasNext = iter.hasNext
    
    override def next() = {
      val pair = (idx, iter.next)
      idx = idx + 1
      pair
    }
  }
  
  def withIndex[A](iter : Iterable[A]) : Iterable[(Int, A)] = 
    new Iterable[(Int, A)](){
      type IA = (Int, A)
      override def elements = new IndexedIterator[A](iter.elements)
      
      override def foreach(body : IA => Unit) = {
        var i = 0
        iter.foreach { x =>
          body(i, x)
          i = i + 1
        }
      }
      
      override def map[B](f : IA => B) : Iterable[B] = {
        var i = 0
        iter.map  { x =>
          val b = f(i, x)
          i = i + 1
          b
        }
      }
      
      override def flatMap[B](f : IA => Iterable[B]) : Iterable[B] = {
        var i = 0
        iter.flatMap { x =>
          val bs = f(i, x)
          i = i + 1
          bs
        }
      }
      
      override def filter(f : IA => Boolean) : Iterable[IA] = {
        val temp = map{x => x}
        temp.filter(f)
      }
    }
}

あれ？id関数ってないの？