Building a segment tree and answering queries

Why it matters

In the last lesson the two simple methods each failed on one half of the work. The plain loop was slow for queries. Prefix sums were slow for updates. This lesson builds the structure that avoids both problems, at least for building and querying. Once you can write these two operations by hand, every variant in this module, such as minimum, maximum and updates, becomes a small change to the same code.

You will build a segment tree over an array of numbers. You will see how it is stored, how it is built, and how a query walks it. You will also see the four programs that check the work, and you will learn the habits that interviewers look for: stating the cost, testing the edge cases, and saying why the size is what it is.

Start with a tree

A tree is a set of nodes connected by links, with one node at the top called the root. Each node can have children below it. A node with no children is a leaf. A binary tree is a tree where each node has at most two children, usually called left and right.

Think of a company's org chart. The chief executive is the root. Each manager is a node with a few people reporting to them. The people with no reports are the leaves. A question such as "how many people are in this department?" can be answered by adding the head count of each sub-team, and that is the same kind of thinking a segment tree uses.

A segment tree is a binary tree in which every node covers a segment, which is a range of the original array. The root covers the whole array. Its two children cover the left half and the right half. Each of those splits again, and so on, until a leaf covers a single element. Every node stores one answer for its segment. For sums, that answer is the total of the segment.

Take the array {2, 4, 1, 7}. The root covers indexes 0 to 3 and stores 14. Its left child covers indexes 0 to 1 and stores 6, because 2 + 4 = 6. Its right child covers indexes 2 to 3 and stores 8, because 1 + 7 = 8. The four leaves store the four values themselves. The picture below shows the whole tree.

A segment tree for the array 2, 4, 1, 7. The root stores 14, its children store 6 and 8, and the four leaves store the values.

Why is this useful? Any range can be built from a few of these stored pieces. The range [1, 3] is the leaf for index 1 plus the node that covers indexes 2 and 3. That is 4 + 8 = 12. We need only two stored answers instead of three loop steps, and on a big array the saving is far larger. The rest of this lesson shows how to find those pieces quickly.

Store the tree in a plain array

We need to store the tree somewhere, and the simplest way in C++ is a vector. The trick is to number the nodes the same way a heap is numbered. The root is node 1. The left child of node i is node 2*i, and the right child is node 2*i + 1. The children of node 2 are nodes 4 and 5. The children of node 3 are nodes 6 and 7. Nothing needs to be stored to say where a node's children are. The position in the array tells you.

The picture below shows the layout for the same four values. Slot 0 is not used, because the numbering starts at 1. Slot 1 holds the root, slot 2 holds the left child, and slot 3 holds the right child.

The segment tree of 2, 4, 1, 7 laid out in an array with slot numbers. Slot 1 holds 14, slots 2 and 3 hold 6 and 8, and slots 4 to 7 hold the values.

Now a question you must answer carefully: how big must the array be? With the layout above, the deepest node can have an index far larger than n. A safe rule is to allocate 4 * n slots. For n = 4 that is 16 slots, and the tree only uses 7 of them. Some slots stay empty, and that is fine. Memory is cheap compared with the mistake of writing past the end of a vector.

Why 4n and not 2n? Take n = 5. The node numbers go up to 15, which is more than 2n = 10, so 2n slots would overflow. The highest number is always below 4n, because the tree has about log2 n levels and the highest number on the deepest level is about 2 raised to that depth. This is the kind of answer interviewers like: you give the rule and the reason, not just the number.

First principle: a function that calls itself

Building and querying the tree are easiest to write with recursion. Recursion means a function that calls itself. It sounds strange the first time, but it matches the shape of the tree exactly. A function asked to handle a segment does one of two things. If the segment is a single element, it answers directly. Otherwise it splits the segment in half, asks each half to handle itself, and combines the two answers.

Every recursive function needs a base case, the situation where it stops calling itself and returns an answer. Here the base case is a segment of length 1, which is a leaf. Without a base case, the calls would never end and the program would crash with a stack overflow message, typically Segmentation fault (core dumped) on Linux.

Each call takes a little space on the call stack. The call stack is the list of functions that are running and waiting for an answer. When a function calls another, the new call is pushed on top. When it returns, it is popped off. The stack holds one frame per active call, and each frame stores the function's arguments and local variables. For a segment tree, a frame stores the node number and the segment's left and right ends.

Step through the calls below to see the stack grow and shrink while the tree is built.

Watch what happens to the stack. The call on the root goes down to the left child, then to its left child, until it reaches a leaf. The leaf returns, and the caller combines the results. The stack only goes as deep as the number of levels, which is about log2 n. That depth is why the operations are O(log n) in time, and it is also why the stack uses only O(log n) extra memory at any moment.

Build: combine the children

The build walks the whole tree once. A call builds one node. If the node covers one element, it stores that element. Otherwise it builds the left child, builds the right child, and then stores their sum in the node.

void build(const vector<long long>& a, int node, int l, int r) {
    if (l == r) { t[node] = a[l]; return; }
    int mid = (l + r) / 2;
    build(a, 2 * node, l, mid);
    build(a, 2 * node + 1, mid + 1, r);
    t[node] = t[2 * node] + t[2 * node + 1];
}

Read the function carefully. The parameters l and r are the ends of the segment this call owns, both included. The midpoint mid splits the segment. The left child covers [l, mid] and the right child covers [mid + 1, r]. These two ranges do not overlap, and together they cover the whole segment. That is the rule that makes the tree correct.

The cost is O(n). Each node is visited once, and there are about 2n nodes in total, because a tree over n leaves has about n - 1 internal nodes. Each internal node does one addition, so the build does about n additions.

Query: three cases at each node

A query asks for the sum of [ql, qr]. It starts at the root and asks each node a question: how does your segment relate to the query range? There are only three answers.

1. No overlap. The node's segment lies completely outside the query. Nothing in it counts, so the node returns the neutral value, 0 for a sum.
2. Full overlap. The node's segment lies completely inside the query. The stored answer is exactly what we need, so the node returns it at once and does not look deeper.
3. Partial overlap. The segment and the query overlap, but the segment sticks out on one side. The node cannot answer alone. It asks both children and adds their answers.

Here is the function. It matches the three cases in the same order.

long long query(int node, int l, int r, int ql, int qr) const {
    if (qr < l || r < ql) return 0;           // no overlap
    if (ql <= l && r <= qr) return t[node];   // full overlap
    int mid = (l + r) / 2;                    // partial overlap
    return query(2 * node, l, mid, ql, qr) + query(2 * node + 1, mid + 1, r, ql, qr);
}

Look at the first test. The segment [l, r] is outside the query when the query ends before the segment starts, which is qr < l, or when the query starts after the segment ends, which is r < ql. Either test alone means no overlap. The second test is full overlap: the query starts before or at the segment start, and ends after or at the segment end.

Does the order of the first two tests matter? Here it does not change the answer. A segment that is fully inside the query can never be fully outside it, so the two tests never both fire. Keep the habit of checking "no overlap" first anyway, because it reads clearly and it is the case that ends most walks early.

How many nodes does a query visit? At each level of the tree, at most two nodes are partial. Every partial node splits into two children, and only the children that overlap the query go on working. The full-overlap nodes stop immediately. So the work is about two nodes per level, over about log2 n levels. That gives O(log n) for every query, and the proof rests on this one fact: a query range has only two ends, and each end causes at most one partial node per level.

Worked example 1: dump the tree

Before you trust a tree, look inside it. The program below builds the tree for {2, 4, 1, 7} and prints every node with its segment and its sum. It walks the tree from the root down, and the indentation shows the depth. This kind of walk is called a tree traversal. The walk here is pre-order: each node is printed before its children.

The visualizer above walks the tree the same way the program below does. Watch the highlighted node move from the root to the leaves, and watch the order in which nodes are printed. Then type the input into the Input box and run the program.

Input, one item per line:

4
2 4 1 7

Expected output:

node 1 [0..3] sum 14
  node 2 [0..1] sum 6
    node 4 [0..0] sum 2
    node 5 [1..1] sum 4
  node 3 [2..3] sum 8
    node 6 [2..2] sum 1
    node 7 [3..3] sum 7

Read the output as a picture of the tree. Node 1 is the root, and it covers [0..3] with sum 14. Its children are nodes 2 and 3, drawn indented one level. Node 2 covers [0..1] with sum 6, and its children are nodes 4 and 5, the two leaves for 2 and 4. Node 3 covers [2..3] with sum 8, and its children are nodes 6 and 7, the leaves for 1 and 7. Notice that the numbers match the picture earlier: the node numbers are the array slots from the layout diagram.

Use this dump whenever a query gives a wrong answer. Compare each printed sum with the value you expect. If a node's sum is wrong, the bug is in the build. If the sums are right but the query is wrong, the bug is in the query.

Worked example 2: answer queries

Now the query itself. The program below reads the array, builds the tree, and then answers queries from the Input box. It is the complete version of the code you have seen in pieces.

Input, one item per line:

4
2 4 1 7
3
1 3
0 3
2 2

Expected output:

12
14
1

Trace the first query, query(1, 3), on the tree for {2, 4, 1, 7}. The call starts at node 1, which covers [0, 3]. That is a partial overlap with [1, 3], so it asks its children. The left child, node 2, covers [0, 1]. That overlaps [1, 3] partly, because index 1 is in both, so node 2 is partial and asks its own children. Node 4 covers [0, 0], which has no overlap, so it returns 0. Node 5 covers [1, 1], which is inside the query, so it returns its stored 4. Back at node 2 the answer is 0 + 4 = 4.

Now the right child of the root, node 3, covers [2, 3]. That is inside the query, so it returns its stored 8 at once. The root adds 4 and 8 and returns 12. The whole query visited five nodes: the root, nodes 2, 4, 5 and 3. It never used the stored 14 at the root, because the root was partial.

The second query, query(0, 3), covers the whole array. The root is a full overlap at once, so the answer is the stored 14. The third query, query(2, 2), is a single element. It visits the root, which is partial, then node 3, which is partial, then node 6, which is full and returns 1, while node 7 has no overlap and returns 0.

Worked example 3: count the visits

Interviewers ask for the cost of a query, and it is better to measure than to guess. The next program is a maximum tree. It is the same design as the sum tree, with two changes. The combining operation is max, and the neutral value is INT_MIN, the smallest possible int, because it never wins a maximum. The program also counts how many nodes each query visits.

Input, one item per line:

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

Expected output:

max 9 (visited 1 nodes)
max 9 (visited 9 nodes)
max 7 (visited 9 nodes)
max 8 (visited 7 nodes)

The first query covers the whole array. The root is a full overlap, so it returns at once and visits 1 node. The second query, [1, 2], has two values, 1 and 9. It visits 9 nodes, and that count is typical for a range that is not aligned with the tree's split points. Two nodes per level are partial, and the rest are the checks that find no overlap or full overlap. The third query, [3, 5], visits 9 nodes too, and returns the largest of 3, 7 and 2, which is 7. The last query, [6, 6], is one element, and it visits 7 nodes, including the walk down from the root.

Do not expect the same count every time. The number depends on where the range sits. What stays true is the bound. A query visits about two partial nodes per level, plus the few nodes that stop the walk, so the count grows like log n, not like n. Try the same program with 1000 elements and you will see the counts grow by a few, not by hundreds.

Notice the neutral value. If the code returned 0 for no overlap instead of INT_MIN, a range of all negative numbers would wrongly return 0. A neutral value must never win the comparison. For a maximum that means the smallest possible value. The common mistakes section shows what goes wrong when you pick the wrong one.

Worked example 4: check against the simple loop

A new tree is not trustworthy until it agrees with a method you already trust. The simplest trusted method is the loop. The next program builds the tree and then checks every possible range [l, r] against a plain loop. It counts the checks and the mismatches. A correct tree prints zero mismatches.

Input, one item per line:

6
3 1 4 1 5 9

Expected output:

checked 21 ranges, mismatches 0

The number 21 is not an accident. For n values, the number of ranges [l, r] with l <= r is n times (n + 1) divided by 2. For n = 6 that is 6 times 7 divided by 2, which is 21. Knowing the count lets you check the test itself: if the program printed a checked count different from 21, the loops would be wrong, not the tree.

This kind of test is called a brute-force check. Brute force means trying every possibility. It is slow, but it is simple enough to be obviously right, so it makes a fine referee for a fast method. Use it on small arrays while you develop, and never on large ones.

What happens inside: the cost in memory

Think about memory for a moment. A tree over n values uses about 4n slots of long long. For n = 100000 that is 400000 slots, or 3.2 megabytes. That is a modest amount for modern machines, which is why a segment tree is practical even for large arrays.

The call stack during a query holds at most about log2 n frames at once, because the calls only go as deep as the tree is tall. For n = 100000 that is about 17 frames. Each frame is small, so the stack costs almost nothing. A program that builds the tree recursively is therefore safe for large arrays. The danger is the other direction: a recursive function that goes as deep as n would be a problem, and a segment tree never does that.

The tree is also a single block of memory with no pointers. Each node is found from its number with one multiplication and one addition, as in the last lesson's vector lookups. There are no separate heap allocations per node, so the structure stays tidy in memory.

Why the depth is log n

The whole cost argument rests on one fact: the tree is short. A tree over n leaves is about log2 n levels tall, and this is easy to see by halving. Start with a segment of length n. Its children have length about n / 2. Their children have length about n / 4. Keep halving, and after k levels each segment has length about n / 2^k. The segment reaches length 1, which is a leaf, when 2^k is about n, so k is about log2 n.

Try it with real numbers. An array of 8 elements splits 8, then 4, then 2, then 1, which is three halvings, and log2 8 is 3. An array of 1000 elements takes about 10 halvings, because 2^10 is 1024. An array of one million takes about 20, because 2^20 is just over one million. Each step down the tree is one more recursive call, so the depth of the call stack is the same number. Twenty calls deep is nothing for a computer, which is why recursion is a safe choice here.

Because the tree is short, a query can afford to visit about two nodes at each level. Two nodes per level times 20 levels is about 40 node visits, even for a million-element array. Compare that with the loop, which visits every element in the range.

A full trace of a second query

The first query trace in worked example 2 covered a range that started at index 1. Here is a second trace, for query(0, 2) on the same tree for {2, 4, 1, 7}. The range is 2 + 4 + 1 = 7, and the trace shows every call in order.

1. Call on node 1, which covers [0..3]. The query [0..2] overlaps it but does not cover it, because index 3 is outside the query. This is a partial overlap, so the call asks its two children.
2. Call on node 2, which covers [0..1]. Both of its indexes are inside [0..2], so this is a full overlap. It returns the stored 6 at once.
3. Call on node 3, which covers [2..3]. Index 2 is inside the query and index 3 is not, so this is a partial overlap. It asks its children.
4. Call on node 6, which covers [2..2]. This is a full overlap and returns the stored 1.
5. Call on node 7, which covers [3..3]. This is outside the query, so it returns 0.
6. Node 3 returns 1 + 0 = 1, and the root returns 6 + 1 = 7.

Five nodes were visited: the root, nodes 2, 3, 6 and 7. Node 2 was answered by its stored value, and node 7 by the neutral value. Nothing below node 2 was visited, which is the point of stored answers. The loop for this query would have added three numbers. The tree added two stored answers and one leaf, because the stored answer for [0..1] already contained 2 + 4.

Writing a segment tree from scratch

When you write one of these trees for the first time, a short checklist keeps you on track. First, decide the combining operation and its neutral value, and write them down. Second, choose the node numbering. Use the root at 1, children at 2*node and 2*node + 1, and allocate 4 * n slots. Third, write the build with the base case for a leaf, the midpoint, the two recursive calls, and the combine. Fourth, write the query with the three cases in the order no overlap, full overlap, partial overlap. Fifth, test the tree with a dump on a small array and a brute-force check on every range. Each step is small, and a bug in one step shows up in the test of the next.

Most of the bugs in this lesson come from skipping one of these steps. Forgetting the neutral value causes wrong answers for ranges outside the query. Forgetting the base case causes infinite recursion. Choosing a different midpoint in the build and the query causes the children to disagree with the query's expectations. The checklist is the cure.

Common mistakes

Too small an array. Allocating 2 * n slots for the recursive layout is a classic error. The numbering can reach about 4n, so the program writes past the end of the vector. Writing past the end of a vector is undefined behaviour. The program may print a wrong answer, or it may crash with Segmentation fault (core dumped). Allocate 4 * n.

Wrong neutral value. Returning 0 for "no overlap" is right for sums and wrong for maximums and minimums. A maximum tree that returns 0 for an empty part gives 0 as the answer for a range of negative numbers.

Wrong root range. The root covers [0, n - 1], not [1, n], for a 0-based array. Calling query(1, 0, n, ...) makes the root think it covers one more element than it does, so the right child reads past the end.

Missing headers. The maximum tree uses INT_MIN. Without #include <climits> g++ reports error: 'INT_MIN' was not declared in this scope, the same kind of error as the one you saw for INT_MAX in the last lesson.

Forgetting a constructor argument. A class that has a constructor taking the array will not build without one. If you write SumTree st; with no argument, g++ reports:

error: no matching function for call to 'SumTree::SumTree()'

The fix is to pass the array, as in SumTree st(a);.

Debugging walkthrough: the answers are all zero

A learner builds the tree for {2, 4, 1, 7} and every query returns 0, including the query over the whole array. The learner runs the dump program and sees zeros in every node. The sums were never written, so the problem is in the build, not in the query.

The learner's constructor looks like this:

SumTree(const vector<long long>& a) : n(a.size()), t(4 * a.size(), 0) {
    // build(a, 1, 0, n - 1);  // this line was commented out while testing
}

The vector is the right size and every slot starts at 0. Nothing ever calls build, so no node is filled in. Queries then read zeros and add zeros, and they return 0. The fix is to restore the call.

The method here is the one to remember. When a tree gives only zeros, print the nodes first. If every node is zero, the build never ran or never wrote. If the nodes are right but a query is wrong, check the query's three cases and the root range n - 1. Print before you guess.

Variations and edge cases

A tree over a single element is a special case. The root is also a leaf, so the build sets t[1] and returns, and the query on [0, 0] returns the root at once. Test it, because many learners' trees break at n = 1.

A tree over an array whose length is not a power of two works without any special code. The midpoint rule splits any length, so the segments are just uneven. Test lengths 5, 6 and 7 as well as 4 and 8.

A query with ql equal to qr is a single element. It always works, but it is a good test of the no-overlap and full-overlap tests together.

A maximum tree needs INT_MIN as the neutral value, and a minimum tree needs INT_MAX. A sum tree needs 0. If you change the operation, you must change the neutral value too.

For a sum over long long values, the tree needs long long nodes, not int nodes. The same overflow rule from the last lesson applies to every node, because the root holds the total of the whole array.

Interview-style questions

Interviewers often ask you to explain the design, not only to write it. Be ready to answer these.

"Why is a query O(log n)?" Answer: the query range has two ends, and each end creates at most one partial node per level. The tree has about log2 n levels, so the partial nodes are about 2 log2 n in number, and the full-overlap nodes stop the walk early.

"How big must the array be?" Answer: at most about 4n, because the tree has about log2 n levels and the numbering of a non-power-of-two tree can skip slots. Allocating 4n always works, and it is cheap.

"Why do you need the neutral value?" Answer: a node outside the query must not change the result, so it returns a value that does not win the combining operation.

"Could you use a loop instead?" Answer: yes, for small or static data. The point of the tree is to make both queries and updates O(log n) at the same time, which the loop cannot do.

Practice

1. Build the tree for {3, 1, 4, 1} by hand and write down all seven node values in order from node 1 to node 7.
2. Write a maximum tree for {3, 1, 4, 1} and check query(0, 3) and query(1, 2) by hand. Which neutral value do you use?
3. Count how many nodes query(0, 0) visits on an array of 8 elements. Use the dump to help. Compare your count with the bound of about 2 log2 n.
4. What goes wrong if build splits with mid = (l + r + 1) / 2 but the query uses (l + r) / 2? Give a small example where the two rules disagree.
5. Change the sum tree in the second program so that it answers the minimum. Name every line you change, and give the new neutral value.
6. Explain why the full-overlap test ql <= l && r <= qr can never be true for a segment that lies outside the query.
7. Add a method int size() const to the sum tree that returns the number of elements. Explain why storing n is needed.
8. Predict the output of the dump program for the array {5}. Then predict the output for {5, 6}. Check both by running the program.
9. The brute-force check reports checked 21. Predict the checked count for n = 10 before you run anything, and explain how you got it.
10. Write a short paragraph, in your own words, for an interviewer explaining why the query only visits about two nodes per level.

What you now know

• A segment tree stores one answer per segment, and the root covers the whole array.
• Nodes live in a vector, numbered from 1, with children at 2*node and 2*node + 1.
• Allocate 4 * n slots, because the numbering can skip some.
• The build combines children, costs O(n), and visits each node once.
• A query checks no overlap, full overlap and partial overlap, in that order.
• A query visits about two nodes per level, so it costs O(log n).
• Check the tree with a dump and with a brute-force loop before you trust it.