SpaceSaving frequent item counter in Python In the code below, I implemented the SpaceSaving frequency estimation algorithm described in this paper. Given a parameter eps, the algorithm finds all elements of a data stream of length n that occur more than n/eps times (with high probability). Here's a screenshot from the paper with the pseudocode:



I would appreciate any feedback on my implementation: style, performance, etc.

import math, heapq

class SpaceSavingCounter:
    def __init__(self, eps):
        self.k = math.ceil(1/eps)
        self.n = 0
        self.counts = dict()
        self.queue = []

    def inc(self, x):
        # increment total elements seen
        self.n += 1

        # x is being watched
        if x in self.counts:
            self.counts[x] += 1

        # x is not being watched
        else:
            # make room for x
            if self.n &gt; self.k:
                while True:
                    count, tstamp, key = self.pop()
                    assert self.counts[key] &gt;= count
                    if self.counts[key] == count:
                        del self.counts[key]
                        break
                    else:
                        self.push(self.counts[key], tstamp, key)
            else:
                count = 0

            # watch x
            self.counts[x] = count + 1
            self.push(count, self.n, x)


    def push(self, count, tstamp, key):
        heapq.heappush(
            self.queue,
            (count, tstamp, key)
        )

    def pop(self):
        return heapq.heappop(self.queue)

def test_SpaceSavingCounter():
    seq = [1,5,3,4,2,7,7,1,3,1,3,1,3,1,3]
    counter = SpaceSavingCounter(1 / 1.9)
    for x in seq:
        counter.inc(x)
    assert counter.counts.keys() == {1,3}
Thanks for putting this code online! Two updates for people who may want to use it:


Correctness: The "if self.n > self.k:" is probably mistaken: for instance, when the sequence starts with k times the same element, we are later stuck with a data structure that contains a single counter. Instead, we can use self.k as representing the number of remaining counters with value 0: specifically, replace this line by "if self.k==0:" and just after the "else:" insert "self.k=self.k-1".
Small optimization: above "def push", replace by "self.push(self.counts[x], self.n, x)", i.e. use the latest value of x since we just updated it