Amortized O(1) and deamortized O(1) functional counters, generalized to arbitrary semigroups

counters.scala

object Amortized {

  /**
   * Amortized constant time counter. Maintains the invariant that
   * `chunks` are of exponentially decreasing sizes. May perform a
   * logarithmic number of `op` reductions per `snoc` (worst case),
   * but a series of `n` snocs will on average require `2n` reductions.
   */
  case class Counter[A](chunks: Vector[A]) {
    def snoc(size: A => Long, op: (A,A) => A)(a: A): Counter[A] = {
      var c2 = chunks :+ a
      while (c2.size > 1 && size(c2.last)*2 > size(c2(c2.length - 2))) {
        c2 = c2.dropRight(2) :+ op(c2(c2.length - 2), c2.last)
      }
      Counter(c2)
    }
  }
  object Counter {
    def empty[A] = Counter[A](Vector.empty)
    def single[A](a: A) = Counter(Vector(a))
  }
}

object Deamortized {

  /**
   * Deamortized worse case constant time counter. Invariant is
   * that `chunks` has exponentially decreasing sizes, and `dirty`
   * has values that are pending `:+` onto `chunks`. At most two
   * `op` reductions are performed on each `snoc`.
   */
  case class Counter[A](chunks: Vector[A], dirty: Vector[A]) {

    def fixup(size: A => Long, op: (A,A) => A): Counter[A] =
      if (dirty.isEmpty) this
      else if (chunks.nonEmpty && size(dirty.head) * 2 > size(chunks.last)) {
        val h2 = op(chunks.last, dirty.head)
        if (chunks.size > 1 && size(h2) * 2 > size(chunks(chunks.size - 2))) {
          Counter(chunks.init, h2 +: dirty.tail)
        }
        else
          Counter(chunks.init :+ h2, dirty.tail)
      }
      else Counter(chunks :+ dirty.head, dirty.tail)

    def snoc(size: A => Long, op: (A,A) => A)(a: A): Counter[A] =
      Counter(chunks, dirty :+ a).fixup(size,op).fixup(size,op)
  }
  object Counter {
    def empty[A] = Counter[A](Vector.empty, Vector.empty)
    def single[A](a: A) = Counter(Vector(a), Vector.empty)
  }
}