Segment tree structure and custom merges

A segment treesegment treeA tree data structure that stores partial answers about ranges, enabling point updates and range queries in logarithmic time.In the glossary stores partial answers about ranges of an array. Each node holds a summary of one contiguous range, and the tree is built bottom-up so that any range query is answered by combining a logarithmic number of nodes.

Why you need a segment tree

Prefix sums answer a range-sum query in constant time after linear preprocessing, but they cannot help once an element changes, since every later prefix sum becomes wrong. Many problems also ask for something a prefix sum was never built to answer at all, such as a range maximum or a range gcd.

A segment tree answers both point updates and range queries in logarithmic time. The cost is a larger constant factor than a prefix sum array, and more careful bookkeeping in the code.

The shape of the tree

Lay the array out left to right, and build a complete binary tree on top of it: each leaf holds one array element, and each internal node holds a summary of its two children's ranges combined.

The tree has n leaves and roughly 2n nodes in total. Numbering nodes from 1, with node 1 as the root, node i's children at 2i and 2i + 1, lets the whole tree live in one flat array. Padding n up to the next power of 2 keeps every node's range a clean power-of-2 size, which simplifies both the build and the query code.

Building bottom-up

Rather than inserting elements one at a time from the root down, place every leaf value directly into positions size through 2 * size - 1. Then walk backward from node size - 1 down to node 1, combining each node's two children into its own summary. Every internal node is combined exactly once, so the whole build costs $O(n)$.

Point updates and range queries

Updating a leaf means writing the new value at its position, then walking up to the root, recombining the two children at each step. That walk touches $O(\log n)$ nodes.

A range query descends from the root instead. At each node, check whether its range sits entirely inside the query range (use it directly and stop), entirely outside it (skip it), or partially overlapping (descend into both children). This also touches $O(\log n)$ nodes in total.

Custom monoids

A monoid is a type paired with an associative way to combine two values: (a combine b) combine c must equal a combine (b combine c), so the order in which sub-ranges are merged never changes the final answer. Plain sums form a monoid, and so does taking the maximum of two values.

Some problems need more than a single number to summarize a range correctly. The classic example is the maximum-subarray-sum query: given a range of possibly negative values, find the largest sum of any contiguous subarray inside it. A single running total is not enough, because the best subarray might not start or end at the range's own boundary. The fix is to summarize each range with four numbers instead of one: the range's total sum, its best prefix sum, its best suffix sum, and its best subarray sum anywhere inside it.

Combining two ranges' summaries needs all four fields. The combined total is just the sum of the two totals. The combined best prefix is either the left range's best prefix on its own, or the left range's whole total extended by the right range's best prefix. The combined best suffix mirrors that on the right side. The combined best subarray is the larger of the two ranges' own best subarrays, or a subarray that crosses the boundary: the left range's best suffix joined to the right range's best prefix.

1import sys2
3
4def combine(a, b):5    # Each node is (total, best_prefix, best_suffix, best_subarray).6    total = a[0] + b[0]7    best_prefix = max(a[1], a[0] + b[1])8    best_suffix = max(b[2], b[0] + a[2])9    best_subarray = max(a[3], b[3], a[2] + b[1])10    return (total, best_prefix, best_suffix, best_subarray)11
12
13def leaf(v):14    return (v, v, v, v)15
16
17def main() -> None:18    input_data = sys.stdin.read().split()19    idx = 020    n = int(input_data[idx])21    idx += 122
23    size = 124    while size < n:25        size *= 226
27    neutral = (0, -10 ** 18, -10 ** 18, -10 ** 18)28    tree = [neutral] * (2 * size)29
30    for i in range(n):31        v = int(input_data[idx])32        idx += 133        tree[size + i] = leaf(v)34
35    for i in range(size - 1, 0, -1):36        tree[i] = combine(tree[2 * i], tree[2 * i + 1])37
38    num_queries = int(input_data[idx])39    idx += 140
41    results = []42    for _ in range(num_queries):43        query_type = int(input_data[idx])44        idx += 145
46        if query_type == 1:47            pos = int(input_data[idx])48            v = int(input_data[idx + 1])49            idx += 250
51            i = size + pos52            tree[i] = leaf(v)53            i //= 254            while i > 0:55                tree[i] = combine(tree[2 * i], tree[2 * i + 1])56                i //= 257        else:58            left = int(input_data[idx])59            right = int(input_data[idx + 1])60            idx += 261
62            # Collect the O(log n) covering nodes in left-to-right order,63            # since the merge is not commutative in its prefix/suffix fields.64            left_stack = []65            right_stack = []66            lo = left + size67            hi = right + size + 168
69            while lo < hi:70                if lo % 2 == 1:71                    left_stack.append(tree[lo])72                    lo += 173                if hi % 2 == 1:74                    hi -= 175                    right_stack.append(tree[hi])76                lo //= 277                hi //= 278
79            acc = None80            for node in left_stack + right_stack[::-1]:81                acc = node if acc is None else combine(acc, node)82
83            results.append(str(acc[3]))84
85    sys.stdout.write("\n".join(results) + "\n")86
87
88if __name__ == "__main__":89    main()

9
-2 1 -3 4 -1 2 1 -5 4
3
2 0 8
1 7 0
2 0 8

6
10

Take the array [-2, 1, -3, 4, -1, 2, 1, -5, 4]. The best contiguous subarray is [4, -1, 2, 1], summing to 6, which the program confirms for the full range. After changing the -5 at position 7 to a 0, the tail of the array from position 3 onward, [4, -1, 2, 1, 0, 4], becomes entirely worth keeping, and its sum of 10 becomes the new best subarray over the full range. Neither answer could come from tracking a single running maximum; both prefix and suffix sums are needed to correctly extend a subarray across the boundary between two combined nodes.

Common mistakes

Forgetting that the flat tree array is effectively rooted at index 1, with the leaves starting at size rather than at 0, is the most common bookkeeping error. Padding n up to a power of 2 but then indexing leaves with the original, unpadded positions produces a tree that runs without error but answers every query against the wrong elements.

Another mistake is combining ranges in the wrong order once the combine operation is not commutative, as with the prefix and suffix fields above. Reversing the order in which a query's covering nodes were meant to be merged produces a plausible-looking but wrong answer, since prefix and suffix sums depend on which side of the range they are attached to.

Using floating-point values with a combine function that a rounding error can make non-associative is a third mistake. If two different merge orders can produce two different floating-point results, the segment tree's whole correctness argument, that merge order never matters, breaks down.

Making it fast on PyPy

Segment tree code spends most of its time on array access, so keep the representation itself simple. Plain lists indexed by integers beat dictionaries by a wide margin, and PyPy's JIT compiler optimizes a small, repeated loop body far better than one with function calls buried inside it. Avoid allocating new objects inside the query or update loop; accumulate the answer in a small number of local variables instead.