Segment tree

A segment tree stores the answer for a range, then for each half of it, then for each quarter, down to single values. Any range can be made from O(log n) of these pieces, and changing one value changes only one piece per level. It works for sums, minimums and almost any other way of combining values.

A tree of ranges

The root covers the whole list. Every node that covers more than one position has two children: the left half of its range and the right half. The leaves cover one position each, and they are the values themselves. Every node stores the answer for its range, here the sum, found from its two children.

In the figure, the array cells are the leaves, and every bar above them is a node over exactly the cells it covers; the two bars just below a bar are its children. With n leaves there are n − 1 other nodes, and about log₂ n levels.

Asking

To find the sum of a range, start at the root and, at each node, compare its range with the asked range:

• no overlap: contribute nothing, and do not go below;
• the node’s range lies completely inside: contribute its stored sum, and do not go below;
• partly inside: ask both children and add their answers.

Why O(log n)? On any level, the nodes that are only partly inside must contain one of the two ends of the asked range, so there are at most two of them. Only those are opened, so each level has at most four nodes visited, and there are about log₂ n levels. Step through “Sum of a[2:7]”.

Changing a value

Changing a[i] changes the sum of every node whose range contains i: exactly one node per level, on the path from the leaf up to the root. Recompute each from its two children, from the bottom up. That is O(log n).

Stored in a list

Like a heap (HEAP-1), the tree fits in a list with no links. A short and fast way, often used in Python, puts the n leaves at positions n to 2n − 1 and the node at position k above its children at 2k and 2k + 1. Position 1 is the root, and position 0 is not used.

class SumTree:
    def __init__(self, a):
        n = self.n = len(a)
        t = self.t = [0] * n + list(a)        # leaves at n .. 2n - 1
        for k in range(n - 1, 0, -1):
            t[k] = t[2 * k] + t[2 * k + 1]

    def set(self, i, v):                      # a[i] = v
        t = self.t
        k = i + self.n
        t[k] = v
        while k > 1:
            k //= 2
            t[k] = t[2 * k] + t[2 * k + 1]

    def total(self, l, r):                    # sum of a[l:r]
        t = self.t
        s = 0
        l += self.n
        r += self.n
        while l < r:
            if l & 1:                         # l is a right child
                s += t[l]
                l += 1
            if r & 1:                         # r - 1 is a left child
                r -= 1
                s += t[r]
            l //= 2
            r //= 2
        return s

This version works from the leaves upwards, and it works for any n, not only powers of two. The two ends of the range move up one level per step. When the left end is a right child, its parent reaches outside the range, so the node is added on its own and the end moves one place right; the right end does the same, mirrored. It visits the same O(log n) nodes as the top-down walk of the figure, without recursion. l & 1 is 1 when l is odd.

Any operation that combines

Nothing above needs subtraction. Replace + by any operation that can join the answers of two neighbouring ranges, and the tree works the same way. Such an operation must be associative: (x ∘ y) ∘ z = x ∘ (y ∘ z), so that the order in which pieces are joined does not matter (∘ stands for the operation). For example:

• minimum or maximum (use min, and start the total at float("inf") instead of 0);
• greatest common divisor;
• the number of times the maximum appears, stored as a pair (maximum, count);
• the largest total of a stretch inside the range, stored as four numbers per node.

If the operation is not commutative (x ∘ y differs from y ∘ x, as for joining strings), keep the left and right pieces apart in total and join them in order. Practice RNGQ-P4 asks for a minimum tree with changes.