Binary Indexed Tree

bit.py

class BIT(object):
    """ Binary Index Tree

    the index starts from 1
    """
    def __init__(self, size):
        self.size = size
        self.xs = [0] * (size + 1)
        self.data = [0] * (size + 1)

    def __setitem__(self, idx, value):
        if idx == 0:
            raise ValueError("idx cannot be 0")
        delta = value - self.data[idx]
        self.add_to(idx, delta)
        self.data[idx] = value

    def add_to(self, idx, delta):
        self.data[idx] += delta
        if idx == 0:
            raise ValueError("idx cannot be 0")
        while idx <= self.size:
            self.xs[idx] += delta
            idx += idx & -idx

    def __getitem__(self, idx):
        return self.data[idx]

    def get_cum(self, idx):
        result = self.xs[idx]
        while idx:
            idx -= idx & -idx
            result += self.xs[idx]
        return result

    def get_range(self, start, end):
        """
        get the cumulative sum in the range of (start, end)
        both ends inclusive
        """
        start = max(1, start)
        end = min(self.size, end)
        return self.get_cum(end) - self.get_cum(start-1)