HuffmanTree.scala

import scala.collection.mutable.HashMap
class HuffmanTree(str: String) {
  abstract class Node {
    val count: Int
  }
  case class Leaf(str:String, count:Int)     extends Node
  case class Branch(left: Node, right: Node) extends Node {
    val count = left.count + right.count
  }

  var queue     = List[Node]()
  var codes     = List[(String, String)]()
  var counts    = new HashMap[String, Int]
  var keys      = new HashMap[String, String]
  var tree:Node = null
  build

  def build = {
    makeCounts
    makeQueue
    makeTree
    makeCode
  }
  def makeCounts = {
    for (i <- 0 until str.length) {
      val key   = str(i).toString
      val count = counts.getOrElse(key, 0) + 1
      counts.update(key, count)
    }
  }
  def makeQueue = {
    counts.foreach(arg =>
      queue = new Leaf(arg._1, arg._2) :: queue
    )
    queue = queue.sortWith((a, b) => a.count < b.count)
  }
  def extractMin:Node = {
    val min = queue.head
    queue   = queue.tail
    return min
  }
  def makeTree = {
    while (queue.size > 1) {
      var branch = new Branch(extractMin, extractMin)
      queue = (branch :: queue).sortWith((a, b) => a.count < b.count)
    }
    tree = queue.head
  }
  def makeCode = {
    codes = getCode(tree, "")
    codes.foreach(code => keys.update(code._1, code._2))
  }
  def getCode(node:Node, code:String): List[(String, String)] = node match {
    case Leaf(str, count)    => List((str, code))
    case Branch(left, right) => getCode(left, code + "0") ::: getCode(right, code + "1")
    case _                   => throw new IllegalArgumentException
  }
  def dump = {
    println(tree)
    println(keys)
  }
}