FenwickTree.py

class Fenwick:
    def __init__(self, length, array = []):
        """
        initialise tree
        """
        self.length = length + 1
        self.tree = [0 for pos in range(self.length)]
        
        if array != []:
            self.array = [0] + array
            #one based indexing
            for i in range(1, self.length):
                self.tree[i] = self.array[i]

            for child in range(1, self.length):
                parent = child + self.findRSB(child)
                if parent < self.length:
                    self.tree[parent] += self.tree[child]  
        else:
            self.array = [0 for i in range(self.length)]
                      
    def update(self, index, val):
        """
        update values of arr[index] to val
        delta = newVal - previousVal
        (zero based indexing in argument)
        """
        
        index += 1
        delta = val - self.array[index]
        self.array[index] = val
        
        pos = index
        while pos < self.length:
            self.tree[pos] += delta
            pos += self.findRSB(pos)
            
    def increment(self, index, delta):
        """
        increment value of arr[index] by delta
        (zero based indexing in argument)
        """
        
        index += 1
        self.array[index] += delta
        
        pos = index
        while pos < self.length:
            self.tree[pos] += delta
            pos += self.findRSB(pos)
            
    def prefixSum(self, index):
        """
        find prefixSum till index
        (one based indexing in argument)
        """
        ans = 0
        while index != 0:
            ans += self.tree[index]
            index -= self.findRSB(index)
            
        return ans
    
    def findRSB(self, i):
        """
        returns the rightmost set bit
        """
        return i & -i
    
class FenwickRangeQueryPointUpdate(Fenwick):
    def query(self, right, left = 0):
        """
        range query from index l to r
        (zero based indexing in argument)
        """
        return self.prefixSum(right + 1) - self.prefixSum(left)
    
class FenwickPointQueryRangeUpdate:
    def __init__(self, length, array = []):
        """
        initialise tree
        """
        self.fTree = Fenwick(length, [])
        self.length_arr = length
               
        
        for i in range(self.length_arr):
            delta = array[i]
            
            self.fTree.increment(i, delta)
            if i + 1 < self.length_arr:
                self.fTree.increment(i + 1, -delta)
                
        
    def range_update(self, left, right, delta):
        """
        increment the values from index l to r by delta
        (zero based indexing in argument)        
        """           
            
        self.fTree.increment(left, delta)
        if right + 1 < self.length_arr:
            self.fTree.increment(right + 1, -delta)
            
    def point_query(self, index):
        """
        return value of index
        (zero based indexing in argument)
        """    
        return self.fTree.prefixSum(index + 1)
        

class FenwickRangeQueryRangeUpdate:
    def __init__(self, length, array = []):
        """
        intialise 2 Fenwick Trees
        """
        self.fTree1 = Fenwick(length, [])
        self.fTree2 = Fenwick(length, [])
        self.length_arr = length
                
        for i in range(self.length_arr):
            delta = array[i]
            
            self.fTree1.increment(i, delta)
            if i + 1 < self.length_arr:
                self.fTree1.increment(i + 1, -delta)

            self.fTree2.increment(i, delta*(i - 1))
            if i + 1 < self.length_arr:
                self.fTree2.increment(i + 1, -delta*i)
            
        
    def increment(self, left, right, delta):
        """
        increment the values from index l to r by delta
        (zero based indexing in argument)
        """
            
        self.fTree1.increment(left, delta)
        if i + 1 < self.length_arr:
            self.fTree1.increment(right + 1, -delta)
        
        self.fTree2.increment(left, delta*(left - 1))
        if i + 1 < self.length_arr:
            self.fTree2.increment(right + 1, -delta*right)
        
    def query(self, right, left):
        """
        range query from index l to r
        (zero based indexing in argument)
        """
        
        return self.prefixSum(right + 1) - self.prefixSum(left)
    
    def prefixSum(self, index):
        """
        find prefixSum till index
        (one based indexing in argument)
        """
        return self.fTree1.prefixSum(index) * (index-1) - self.fTree2.prefixSum(index)