牛の一歩も、一歩は一歩

日々の関心、備忘録などをまとめる.

binary indexed treeのpython実装

AtCoder Library Practice Contest - AtCoderでbinary indexed treeなるものを知ったため, http://hos.ac/slides/20140319_bit.pdf: titleを元に実装(ほとんどまるパクリ). 以下コード.

class BinaryIndexedTree():
    """
    各関数のaは全て引数の値になっている前提.(0~n-1)
    Attributes:
        n: データ数
        data: データ
    """
    def __init__(self, n, data=None):
        self.n=n
        self.data=[0 for _ in range(n)]
        if not data is None:
            for i, a in enumerate(data):
                self.add(i, a)

    def add(self, a, w):
        x=a
        while x<self.n:
            self.data[x]+=w
            x=x|(x+1)

    def sum(self, a):
        if a<0:
            return 0
        ret=0
        x=a
        while x>=0:
            ret+=self.data[x]
            x=(x&(x+1))-1
        return ret