Segment & Fenwick Trees

Range queries with updates in O(log n), counting with coordinate compression, and range updates.

Prefix sums answer range-sum queries in O(1), but a single update forces an O(n) rebuild. When a problem mixes range queries with updates, you need a structure that does both in O(log n). The two standard choices are the Fenwick tree (binary indexed tree) and the segment tree.

These come up less often than heaps or graphs, but they're the intended solution to a recognizable family of hard problems, and they're common at trading firms, in online assessments and in some Google interviews.

Signals

"Update an element, then query the sum (or min, max) of a range", repeated many times
"Count the elements smaller than x that appear before (or after) each position"
Counting inversions or "reverse pairs" with a twist
Many interval updates followed by point or range queries

Rule of thumb: use a Fenwick tree for sums and counts. Use a segment tree when the operation isn't a sum (min, max, gcd) or when you need range updates.

Fenwick tree (binary indexed tree)

A Fenwick tree stores partial sums in a 1-indexed array. Index i is responsible for a block of length i & -i (its lowest set bit). Adding the lowest set bit walks up to the blocks that contain i; subtracting it walks down through the blocks that make up a prefix.

class Fenwick:
    def __init__(self, n):
        self.tree = [0] * (n + 1)          # 1-indexed

    def add(self, i, delta):               # a[i] += delta, i is 0-indexed
        i += 1
        while i < len(self.tree):
            self.tree[i] += delta
            i += i & -i

    def prefix(self, i):                   # sum of a[0..i-1]
        total = 0
        while i > 0:
            total += self.tree[i]
            i -= i & -i
        return total

    def range_sum(self, lo, hi):           # sum of a[lo..hi], inclusive
        return self.prefix(hi + 1) - self.prefix(lo)


class NumArray:                            # Range Sum Query - Mutable
    def __init__(self, nums):
        self.nums = nums[:]
        self.fw = Fenwick(len(nums))
        for i, x in enumerate(nums):
            self.fw.add(i, x)

    def update(self, i, val):
        self.fw.add(i, val - self.nums[i])
        self.nums[i] = val

    def sum_range(self, lo, hi):
        return self.fw.range_sum(lo, hi)


arr = NumArray([1, 3, 5])
print(arr.sum_range(0, 2))   # 9
arr.update(1, 2)
print(arr.sum_range(0, 2))   # 8

Building by calling add n times is O(n log n), which is fine in an interview.

Counting with a Fenwick tree

Many "count smaller" problems become: walk the array, and for each value ask "how many values seen so far are less than x?". A Fenwick tree indexed by value answers that. If values are large or negative, compress coordinates first: replace each value with its rank in the sorted list of distinct values.

def count_smaller(nums):
    ranks = {v: i for i, v in enumerate(sorted(set(nums)))}
    fw = Fenwick(len(ranks))
    out = []
    for x in reversed(nums):
        r = ranks[x]
        out.append(fw.prefix(r))          # values already seen with rank < r
        fw.add(r, 1)
    return out[::-1]

print(count_smaller([5, 2, 6, 1]))        # [2, 1, 1, 0]

The same idea counts inversions in O(n log n), as an alternative to the merge-sort method in Sorting & Quickselect.

Segment tree

A segment tree is a binary tree over the array. Each node stores the answer (sum, min, max…) for a range, and its two children split that range in half. A query combines O(log n) nodes that exactly cover the requested range; an update changes one leaf and the nodes above it.

The iterative version below is short and works for any associative operation with an identity value.

class SegmentTree:
    def __init__(self, nums, op=min, identity=float("inf")):
        self.n = len(nums)
        self.op, self.identity = op, identity
        self.tree = [identity] * (2 * self.n)
        self.tree[self.n:] = nums                     # leaves
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = op(self.tree[2 * i], self.tree[2 * i + 1])

    def update(self, i, val):
        i += self.n
        self.tree[i] = val
        while i > 1:
            i //= 2
            self.tree[i] = self.op(self.tree[2 * i], self.tree[2 * i + 1])

    def query(self, lo, hi):                          # op over nums[lo..hi], inclusive
        res = self.identity
        lo += self.n
        hi += self.n + 1
        while lo < hi:
            if lo & 1:
                res = self.op(res, self.tree[lo])
                lo += 1
            if hi & 1:
                hi -= 1
                res = self.op(res, self.tree[hi])
            lo //= 2
            hi //= 2
        return res


st = SegmentTree([5, 2, 8, 6, 3, 7])
print(st.query(0, 5), st.query(2, 4))   # 2 3
st.update(1, 9)
print(st.query(0, 2))                   # 5

Pass op=max, identity=float("-inf") for range max, or op=lambda a, b: a + b, identity=0 for sums.

Range updates and lazy propagation

If the problem says "add v to every element in [l, r]" many times:

Range update, point query: use a Fenwick tree over a difference array. Add v at l and -v at r + 1; the value at i is the prefix sum up to i.
Range update, range query: a recursive segment tree with lazy propagation. Each node stores a pending update that it pushes to its children only when a later operation needs to go below it.

Lazy propagation is long and easy to get wrong under pressure. In an interview, first check whether a simpler idea works: a difference array if all updates come before all queries, a sweep line over sorted events, or a balanced structure such as a sorted list.

fw = Fenwick(6)

def range_add(lo, hi, v):
    fw.add(lo, v)
    if hi + 1 < 6:
        fw.add(hi + 1, -v)

def point_value(i):
    return fw.prefix(i + 1)

range_add(1, 3, 5)
range_add(2, 5, 1)
print([point_value(i) for i in range(6)])   # [0, 5, 6, 6, 1, 1]

[!INTERVIEW] Lead with the simpler solution Start with the brute force or prefix sums and say why it's too slow: "Each update costs O(n), and there are 10⁵ updates." Then introduce the tree. Interviewers care that you recognize when you need it at least as much as the implementation.

Common mistakes

Mixing 0- and 1-indexing in a Fenwick tree. Keep the shift inside the class.
Forgetting coordinate compression when values go up to 10⁹ or are negative.
Using the wrong identity: 0 for sums, +∞ for min, −∞ for max.
Reaching for lazy propagation when the updates and queries don't interleave.

[!TIP] Practice Count of Smaller Numbers After Self · Reverse Pairs · Range Sum Query - Immutable · My Calendar III