huffman encoding

gistfile1.scala

object Huffman extends App {
  sealed trait Tree { val weight: BigDecimal }
  case class Leaf(override val weight: BigDecimal) extends Tree
  case class Node(override val weight: BigDecimal, lhs: Tree, rhs: Tree) extends Tree
  
  def mkHuffmanTree(input: List[BigDecimal]) = {
    var xs: List[Tree] = input.sorted.map(x => Leaf(x))
    Iterator.iterate(xs){
      xs0 => {
        val xs1 = xs0.sortWith(_.weight < _.weight)
        val weight = xs1(0).weight + xs1(1).weight
        Node(weight, xs1(0), xs1(1)) :: xs1.drop(2)
      }
    }.dropWhile(_.length != 1).next.head
  }

  def encode(tree: Tree) = {
    var xs = List[(BigDecimal, String)]()
    def rec(tree: Tree, sign: String): Unit = 
      tree match {
    case Leaf(w) =>
          xs ::= (w, sign)
    case Node(_, l, r) =>
          rec(l, "1"+sign)
          rec(r, "0"+sign)
      }
    rec(tree, "")
    xs
  }

  def efficiency(codes: List[(BigDecimal, String)]) =
    codes.map(x => x._1 * x._2.length).sum

  def craft(codes: List[(BigDecimal, String)]) =
    codes.map(x => 1 / math.pow(2, x._2.length)).sum

  val ws = List[BigDecimal](0.4, 0.15, 0.14, 0.12, 0.10, 0.09)
  
  def run = efficiency(encode(mkHuffmanTree(ws)))

  println(run)
}