Fenwick Tree (or Binary Indexed Tree) Python Implementation, 308. Range Sum Query 2D - Mutable

 

Fenwick Tree (or Binary Indexed Tree) is very efficient in dealing with range queries, especially when elements get updated occasionally.

For how it works, refer to the below video. He elaborates on the principle of the tree in a very clean way.

https://www.youtube.com/watch?v=CWDQJGaN1gY&t=160s&ab_channel=TusharRoy-CodingMadeSimple 

Below is the python implementation of it.

class Fenwick:
    """
    Python implementation of
    Fenwick Tree (or Binary Indexed Tree)
    # Useful in handling range queries
    # Update: O(logN)
    # query: O(logN)
    # buildTree: O(NlogN)
    """
    def __init__(self, nums):
        """
        Initialization
        """
        self.N = len(nums)
        self.nums = [0] * self.N
        self.tree = [0] * (self.N+1)
        for i, num in enumerate(nums):
            self.update(i, num)

    def update(self, i: int, val: int):
        """
        Updates value at index i with val
        # val: new value
        """
        diff = val - self.nums[i]
        self.nums[i] = val
        i += 1
        while i <= self.N:
            self.tree[i] += diff
            i += i & (-i)

    def query(self, i, j):
        """
        returns sum of nums from i to j (both i & j are included)
        """
        ret = 0
        j = j + 1
        while j > 0:
            ret += self.tree[j]
            j -= j & (-j)
        while i > 0:
            ret -= self.tree[i]
            i -= i & (-i)
        return ret

 

1D Query

https://leetcode.com/problems/range-sum-query-mutable/description/

Range Sum Query - Mutable - LeetCode

class Fenwick:
    """
    Python implementation of
    Fenwick Tree (or Binary Indexed Tree)
    # Useful in handling range queries
    # Update: O(logN)
    # query: O(logN)
    # buildTree: O(NlogN)
    """
    def __init__(self, nums):
        """
        Initialization
        """
        self.N = len(nums)
        self.nums = [0] * self.N
        self.tree = [0] * (self.N+1)
        for i, num in enumerate(nums):
            self.update(i, num)

    def update(self, i: int, val: int):
        """
        Updates value at index i with val
        # val: new value
        """
        diff = val - self.nums[i]
        self.nums[i] = val
        i += 1
        while i <= self.N:
            self.tree[i] += diff
            i += i & (-i)

    def query(self, i, j):
        """
        returns sum of nums from i to j (both i & j are included)
        """
        ret = 0
        j = j + 1
        while j > 0:
            ret += self.tree[j]
            j -= j & (-j)
        while i > 0:
            ret -= self.tree[i]
            i -= i & (-i)
        return ret


class NumArray:

    def __init__(self, nums: List[int]):
        self.fenwick = Fenwick(nums)

    def update(self, index: int, val: int) -> None:
        self.fenwick.update(index, val)

    def sumRange(self, left: int, right: int) -> int:
        return self.fenwick.query(left, right)
        


# Your NumArray object will be instantiated and called as such:
# obj = NumArray(nums)
# obj.update(index,val)
# param_2 = obj.sumRange(left,right)

 

 

2D Query can be handled as well!

LeetCode: 308. Range Sum Query 2D - Mutable

https://leetcode.com/problems/range-sum-query-2d-mutable/description/

Range Sum Query 2D - Mutable - LeetCode

class NumMatrix:

    def __init__(self, matrix: List[List[int]]):
        self.M, self.N = len(matrix), len(matrix[0])
        self.matrix = [[0] * (self.N) for _ in range(self.M)]
        self.tree = [[0] * (self.N+1) for _ in range(self.M+1)]
        for y in range(self.M):
            for x in range(self.N):
                self.update(y, x, matrix[y][x])

    def update(self, row: int, col: int, val: int) -> None:
        delta = val - self.matrix[row][col]
        self.matrix[row][col] = val
        i = row + 1
        while i <= self.M:
            j = col + 1
            while j <= self.N:
                self.tree[i][j] += delta
                j += j & (-j)
            i += i & (-i)

    def sum(self, row, col):
        res, i = 0, row+1
        while i:
            j = col+1
            while j:
                res += self.tree[i][j]
                j -= (j & -j)
            i -= (i & -i)
        return res

    def sumRegion(self, row1: int, col1: int, row2: int, col2: int) -> int:
        return self.sum(row2, col2) + self.sum(row1-1, col1-1) - self.sum(row1-1, col2) - self.sum(row2, col1-1)

            


# Your NumMatrix object will be instantiated and called as such:
# obj = NumMatrix(matrix)
# obj.update(row,col,val)
# param_2 = obj.sumRegion(row1,col1,row2,col2)