case classを使ってtreeを

30分プログラム、その375。Programming in Scala, Third Editionで、とうとうバリアント型ことcase classが登場したのでTreeを作ってみた。
本当は二分木が作りたかったのだけれども、型パラメータTに条件をつける方法が分からなかったので、深さを測る関数ぐらいしか作れなかった。たぶんTraitが関係してると思うんだけどな。

使い方

scala> val tree = Branch(42,Branch(10,Leaf(),Leaf()),Leaf())
tree: Branch[Int] = Branch(42,Branch(10,Leaf(),Leaf()),Leaf())

scala> depth2(tree)
res01: Int = 2

ソースコード

sealed abstract class Tree[T]
case class Leaf[T]()  extends Tree[T]
case class Branch[T](value : T ,left : Tree[T],right : Tree[T]) extends Tree[T]

def depth[T](tree : Tree[T]) : Int =
  tree match {
    case Leaf() => 
      0
    case Branch(_,left,right) => {
      val l = depth(left)
      val r = depth(right)
      (l max r) + 1
    }
  }

def map[T,S](f : T => S,tree : Tree[T]) : Tree[S] =
  tree match {
    case Leaf() => 
      Leaf()
    case Branch(v,left,right) =>
      Branch(f(v),map(f,left),map(f,right))
  }

def fold[S,T](f : (T,S,S) => S,init : S,tree : Tree[T]) : S =
  tree match {
    case Leaf() =>
      init
    case Branch(v,left,right) => 
      f(v,fold(f,init,left),fold(f,init,right))
  }

def depth2[T](tree : Tree[T]) : Int =
  fold(((v : T ,x : Int,y : Int)=>x+y+1),0,tree)