gistfile1.scala

import scala.collection.Traversable
import scala.collection.TraversableLike
import scala.collection.mutable.Builder
import scala.collection.mutable.ListBuffer
import collection.generic.CanBuildFrom
   

class Bag[T] private ( private val bag: Map[T,Int], private val total: Int ) 
extends Traversable[(T,Int)] 
with TraversableLike[(T,Int),Bag[T]] 
with ( T => Int ) 
{
  
  def foreach[U]( f: ((T,Int)) => U ) =
    bag foreach f

  def foreach[U]( f: (T,Int) => U ) = bag foreach f.tupled 

  def apply( t: T ): Int = bag.get( t ).getOrElse( 0 )

  override def newBuilder = Bag.newBuilder[T]

  override def toString(): String =
    "Bag(" + bag.mkString(",") + ")"

  def +( kv: (T,Int) ): Bag[T] =
    new Bag[T]( bag + kv, total + kv._2 )

  def -( t: T ): Bag[T] = 
    new Bag[T]( bag - t, total - bag( t ) )

  def get( t: T ): Option[Int] = {
    if ( bag.contains( t ) ) {
      Some( bag( t ) )
    } else {
      Some( 0 )
    }
  }
  
  def freq( t: T ): Double = 1.0*bag( t )/total
  
}


object Bag {

  def fromSeq[T]( ts: Seq[(T,Int)] ) = {
    val (bag, total ) =
      ts.foldLeft( Map[T,Int]() -> 0 ) { ( previous, kv ) =>
        val (prevMap,prevTotal) = previous
        val (key,count) = kv
        val nextMap = if( prevMap contains key ) 
          prevMap + ( key -> ( count + prevMap(key) ) )
        else
          prevMap + ( key -> count )
        val nextTotal = prevTotal + count
        ( nextMap, nextTotal )
      }
    new Bag( bag, total )
  }
    
  def apply[T]() = new Bag[T]( Map[T,Int](), 0 )
  
  def apply[T]( map: Map[T,Int] ) = {
    val total = map.foldLeft( 0 )( _ + _._2 )
    new Bag[T]( map, total )
  }

  def newBuilder[T] = ListBuffer[(T,Int)]() mapResult fromSeq
  
  implicit def canBuildFrom[T] = 
    new CanBuildFrom[Traversable[(T,Int)], (T,Int), Bag[T]] {
      def apply() = newBuilder
      def apply(from: Traversable[(T,Int)]) = newBuilder
    }

}