Segment Tree and Fenwick Tree: Range Query Data Structures

Segment Tree and Fenwick Tree

These solve range query problems in O(log n) per query and update, versus O(n) naive. Essential for competitive programming and senior interviews.

class SegmentTree:
    def __init__(self, arr):
        self.n    = len(arr)
        self.tree = [0] * (4 * self.n)
        self._build(arr, 0, 0, self.n-1)

    def _build(self, arr, node, start, end):
        if start == end:
            self.tree[node] = arr[start]
        else:
            mid = (start+end) // 2
            self._build(arr, 2*node+1, start, mid)
            self._build(arr, 2*node+2, mid+1, end)
            self.tree[node] = self.tree[2*node+1] + self.tree[2*node+2]

    def query(self, l, r, node=0, start=None, end=None):
        if start is None: start, end = 0, self.n-1
        if r < start or end < l: return 0
        if l <= start and end <= r: return self.tree[node]
        mid = (start+end) // 2
        return (self.query(l, r, 2*node+1, start, mid) +
                self.query(l, r, 2*node+2, mid+1, end))

arr = [1, 3, 5, 7, 9, 11]
st  = SegmentTree(arr)
print(st.query(1, 4))  # 24 (3+5+7+9)

Fenwick Tree (Binary Indexed Tree)

class FenwickTree:
    def __init__(self, n):
        self.tree = [0] * (n+1)
        self.n    = n

    def update(self, i, delta):
        while i <= self.n:
            self.tree[i] += delta
            i += i & (-i)

    def query(self, i):
        total = 0
        while i > 0:
            total += self.tree[i]
            i -= i & (-i)
        return total

    def range_query(self, l, r):
        return self.query(r) - self.query(l-1)