New counter

[pathikrit]

import scala.collection.mutable

class Counter[A](allowNegativeCount: Boolean) extends mutable.Map[A, Int] {
  private[this] val counter = mutable.Map.empty[A, Int] withDefaultValue 0

  def +=(key: A): this.type = this += (key -> 1)

  override def -=(key: A) = this += (key -> -1)

  def ++=(xs: TraversableOnce[A]): this.type = {
    xs foreach +=
    this
  }

  override def +=(kv: (A, Int)) = {
    val (k, v) = kv
    counter(k) += v
    if (counter(k) == 0 || (!allowNegativeCount && counter(k) < 0)) counter.remove(k)
    this
  }

  override def get(key: A) = counter.get(key)

  override def iterator = counter.iterator
}

object Counter {
  def empty[A](allowNegativeCount: Boolean): Counter[A] = new Counter[A](allowNegativeCount)
  def apply[A](xs: TraversableOnce[A], allowNegativeCount: Boolean): Counter[A] = empty[A](allowNegativeCount) ++= xs
}