Fenwick Range Tree (Binary Indexed Tree) on top of Redis

fenwick.py

import sys
from calendar import timegm
from collections import Counter, defaultdict
from datetime import datetime, timedelta

import redis


class DictBackend(object):
    def __init__(self, maximum=None):
        if not maximum:
            maximum = 7
        self.max = maximum
        self.clear()

    def write(self, pairs):
        for ix, val in pairs:
            self._tree[ix] += val

    def read(self, ixs):
        return [self._tree[ix] for ix in ixs]

    def clear(self):
        self._tree = defaultdict(int)


class RedisHashBackend(object):
    def __init__(self, key, maximum=None, client=None):
        if not maximum:
            maximum = sys.maxint >> 1
        self.max = maximum
        self.key = key
        if not client:
            client = redis.Redis()
        self.client = client

    def write(self, pairs):
        with self.client.pipeline() as pipe:
            for ix, val in pairs:
                pipe.hincrby(self.key, str(ix), val)
            pipe.execute()

    def read(self, ixs):
        if not ixs:
            return []
        vals = self.client.hmget(self.key, *map(str, ixs))
        return [int(v or 0) for v in vals]

    def clear(self):
        self.client.delete(self.key)


class BinaryIndexedTree(object):
    def __init__(self, backend):
        self._backend = backend

    def __getitem__(self, ix):
        if isinstance(ix, slice):
            start = max(ix.start, ix.stop)
            stop = min(ix.start, ix.stop)
            return self.cumfreq(start) - self.cumfreq(stop)
        else:
            return self.get(ix)

    def get(self, ix):
        ixs = [ix]
        if ix > 0:
            parent = ix & (ix - 1)
            ix -= 1
            while parent != ix:
                ixs.append(ix)
                ix = ix & (ix - 1)
        vals = self.backend.read(ixs)
        return vals[0] - sum(vals[1:])

    def cumfreq(self, ix):
        ixs = []
        while ix > 0:
            ixs.append(ix)
            ix = ix & (ix - 1)
        return sum(self._backend.read(ixs))

    def update(self, iterable):
        counts = Counter(iterable)
        for ix, val in counts.iteritems():
            if not 0 < ix <= self._backend.max:
                raise ValueError("value %r not in range (0, %r]" % (
                    ix, self._backend.max))
            pairs = []
            while ix <= self._backend.max:
                pairs.append((ix, val))
                ix += ix & -ix
            self._backend.write(pairs)

    def clear(self):
        self._backend.clear()


class TimeTree(BinaryIndexedTree):
    def __init__(self, backend_cls, since=None, until=None, resolution=None):
        if not since:
            since = datetime.min
        if not until:
            until = datetime.max
        if not resolution:
            resolution = timedelta(seconds=60)
        self.resolution = resolution
        self.since = since
        self.until = until
        self._offset = -timegm(self.since.timetuple())
        maximum = self.scale(self.until)
        backend = backend_cls(maximum=maximum)
        super(self.__class__, self).__init__(backend)

    def scale(self, dt):
        stamp = timegm(dt.timetuple()) + self._offset
        total_secs = self.resolution.total_seconds()
        stamp = int((stamp - (stamp % total_secs)) // total_secs)
        return stamp + 1

    def cumfreq(self, dt):
        ix = self.scale(dt)
        return super(self.__class__, self).cumfreq(ix)

    def get(self, dt):
        ix = self.scale(dt)
        return super(self.__class__, self).get(ix)

    def update(self, iterable):
        counter = Counter(self.scale(v) for v in iterable)
        super(self.__class__, self).update(dict(counter.iteritems()))

Looks good, sadly I am not using python in the project that needs it.

BTW, does anybody happen to have any implementation in Lua?