Building and querying a segment tree in Python

Why it matters

Building and querying are the heart of the structure. Once you can write them, the minimum version, the maximum version and the lazy version are small changes to the same code. Interviewers also ask about them directly, because the three-case search is a clean way to see whether you understand divide and conquer.

In the last lesson you saw the tree as a picture and learned how its nodes are numbered. This lesson turns the picture into code. You will write the build, write the query, trace the query by hand and by program, and measure how many nodes a query visits on a very long list.

Recalling the tree

Each node covers a segment [l..r] and stores the sum of that segment. The root is node 1 and covers the whole list. The left child of node i is 2 * i, and the right child is 2 * i + 1. A node whose segment holds one item is a leaf. The list that stores the tree has 4 * n slots, so every node has room no matter how the halves split.

The tree is built once, before any query runs. The build fills every node, so afterwards each node holds the right sum. The query then reads the stored sums. It does not add items one by one.

Building

The build is recursive. Imagine the method being given a segment [l, r] and a node number. The method has two jobs. It must fill the node for that segment, and it must build the smaller segments below it. The method works like this.

1. If l == r, the segment is one item. Store a[l] in the node and stop.
2. Otherwise find mid = (l + r) // 2. Build the left child for [l, mid] at node 2 * node. Build the right child for [mid + 1, r] at node 2 * node + 1.
3. Set the node to the sum of its two children.

The children are built first, because step 3 needs their values. That is the rule of the build. A parent is always finished after its children. Each node is visited exactly once, and there are about 2n nodes, so the build costs O(n). The operation is linear, which means the total time grows in step with the list length.

Python has no separate class for a tree node. The class below keeps the list and the length together, and it uses two helper methods that call themselves. The methods start with an underscore, which tells a reader that they are helpers for the class and not part of its public use. Python does not enforce this, but the convention is common.

Here is a complete program. It builds the tree for [1, 3, 5, 7, 9, 11] and runs four queries. Press Run and you will see the output.

The output is:

15
36
9
32

The four answers are easy to check by hand. 1..3 covers 3, 5 and 7, which add to 15. The whole list adds to 36. Index 4 alone holds 9. Indexes 2 to 5 hold 5, 7, 9 and 11, which add to 32.

Notice the shape of the code. The public method query(l, r) only starts the search at the root. The private method does the real work, and it takes the node number and the node's range as extra parameters. The public method hides the bookkeeping, and the private method carries the state.

Querying: the three cases

A query asks about a range [l..r]. The search starts at the root, which covers [0..n-1], and compares the node's segment with the query range. Each visit falls into exactly one of three cases.

The first two cases are the stopping cases. They answer at once and never look deeper. The third case is the only one that recurses. Look at the test for no overlap. The segment is outside the range when it ends before the query starts, or when it starts after the query ends. Those two conditions are written r < nl and nr < l. The test for full overlap is l <= nl and nr <= r, which says the whole segment fits inside the query.

For the query query(1, 3) on [1, 3, 5, 7, 9, 11], the items the answer needs are at indexes 1, 2 and 3. The leaf for index 0 lies outside, so it returns 0. The leaves for indexes 1, 2 and 3 lie inside, so their stored values are used. The nodes that cover [0..2], [3..5] and [0..1] are partly inside, so the search must look at their children.

Why the query is fast

Here is the key argument. At any level of the tree, the query range can partly cover at most two nodes. The reason is simple. The range has two ends, and each end can cut at most one node on that level. Every other node on that level is either fully inside or fully outside. A fully inside node stops the search at once, and a fully outside node stops it too.

A partly covered node makes two calls, one for each child. So each level contributes at most four calls to the search. The tree has about log2(n) levels, so a query makes at most about 4 * log2(n) calls. For a list of one million items, log2(n) is about 20, so the search makes at most about 80 calls. That is why the query is O(log n).

Think of a company that has to report the sales for teams 3 to 5. It does not add up every person. It takes the department totals that fit inside the range, and it looks deeper only at the departments at the edges of the range. Only the edges need detail. The middle of the range is already summarised in the totals.

The query code, line by line

Read the _query method slowly. It is short, but each line matters.

The first if is the no-overlap test. If the node's right end is before the query's left end, or the node's left end is after the query's right end, the node contributes nothing. The method returns 0. For a sum, 0 is the right answer because adding 0 changes nothing. For a minimum it would be wrong, and the next lesson explains why.

The second if is the full-overlap test. If the node's segment fits entirely inside the query, the stored sum is the exact answer for that part, so the method returns it.

The last three lines handle partial overlap. They compute the midpoint, then ask the left child and the right child for their answers. The two answers are added. The left child covers [nl..mid] and the right child covers [mid+1..nr], so every item in the segment is counted once.

The order of the tests matters. The no-overlap test must come first, because a segment outside the range is also not fully inside it. The full-overlap test must come before the recursion, because otherwise the search would go deeper than needed and still get the right answer more slowly.

A query traced step by step

Now watch a single query closely. The program below builds the same tree and answers query(1, 3). It prints one line each time a node is visited, saying which case applies. The only change from the program above is the printed line in _query.

The output is:

node 1 [0..5]: partial overlap, ask both children
node 2 [0..2]: partial overlap, ask both children
node 4 [0..1]: partial overlap, ask both children
node 8 [0..0]: no overlap, return 0
node 9 [1..1]: full overlap, return 3
node 5 [2..2]: full overlap, return 5
node 3 [3..5]: partial overlap, ask both children
node 6 [3..4]: partial overlap, ask both children
node 12 [3..3]: full overlap, return 7
node 13 [4..4]: no overlap, return 0
node 7 [5..5]: no overlap, return 0
sum of 1..3 = 15

Read the trace in order. The root is partly covered, so the search asks for node 2 and node 3. Node 2 covers [0..2], which is also partly covered, so it asks for node 4 and node 5. Node 4 covers [0..1], which is partly covered, so it asks for node 8 and node 9. Node 8 covers index 0, which is outside the range, so it returns 0. Node 9 covers index 1, which is inside, so it returns 3. Back at node 4, the answer is 0 + 3 = 3. Node 5 covers index 2, inside the range, so it returns 5. Node 2 now has 3 + 5 = 8.

The search then moves to node 3, which covers [3..5]. It is partly covered, so it asks node 6 and node 7. Node 6 covers [3..4] and is partly covered. Node 12 covers index 3, inside, and returns 7. Node 13 covers index 4, outside, and returns 0. Node 7 covers index 5, outside, and returns 0. Node 3 has 7 + 0 + 0 = 7. The root adds 8 and 7 to get 15.

Notice that node 7 is entirely outside the range, so the search stopped there without looking deeper. That early stop is what keeps the work small.

The trace also shows why the answer is right. Every item from index 1 to index 3 was added exactly once. Index 1 came from node 9, index 2 from node 5, and index 3 from node 12. No node was counted twice, because each leaf is reached through one path and the three cases never add the same item in two places.

Counting visits on a long list

The trace above is useful for a small list. For a long list you need a number instead. The next program builds a tree over the values 1, 2, 3, and so on up to n, then counts how many calls a query makes. It prints the count so you can see how little work the search does compared with the length of the list.

The program below uses a list of one million items and the query from index 10 to index 999990. The list holds the values 1 to 1000000, so index i holds the value i + 1.

The output is:

items: 1000000
sum of 10..999990 = 499991499981
nodes visited: 75

The query covers almost the whole list, yet it visits only 75 nodes. A loop would have added nearly one million numbers. The sum is also worth checking by hand. The values at indexes 10 to 999990 are 11 to 999991. The sum from 1 to 999991 is 999991 * 999992 // 2, which is 499991500036. Take away the first ten values, 1 to 10, which add to 55, and you get 499991499981. The tree gives the same answer.

Python integers grow as large as they need to be, so this sum cannot overflow. A fixed-size integer in some other language would wrap around or fail when the total gets too large. Python's int has no such limit, which is one place where Python is simpler than a language such as Java. The price is that large integers take a little more memory and time than small ones.

Changing the range in the last line changes the number of visits. A range that starts at 0 and ends at 999999 covers everything, and it visits only one node, because the root is fully inside and returns at once. A query on a single item visits about 40 nodes in this tree (41 for most positions). That is one path of about 20 nodes, plus one sibling for each level that the search checks and returns from at once. All of these stay small compared with the list.

Verifying the tree against a loop

A segment tree is easy to get almost right. The best defence is to check it against a simple loop on many ranges. The loop is slow but obviously correct, so it makes a good reference. The next program builds trees for every size from 1 to 12, and for each tree it checks every possible range. It counts the checks and the passes.

The program has no input. Press Run and you will see the result.

The output is:

ranges tested: 364
passed: 364

The count 364 is not random. For a list of size s, the number of ranges is s * (s + 1) / 2, and the sizes from 1 to 12 add to 364. Every range passed. This test covers the awkward small cases: one-item lists, ranges of one item, and ranges that touch the ends. Those are the cases where off-by-one errors hide. A test that only checks the big example may miss them.

The reference uses a slice, a[l:r + 1]. A slice stops before its end, so r + 1 is what makes the range include index r. The random numbers come from random.Random(42). A seed fixes the sequence, so the same program prints the same numbers every time. This matters for tests, because a failing test must fail the same way on the next run. The module-level random.seed(42) does the same job, but a separate Random object keeps the test from changing the state of other code that uses random.

What happens inside Python

The query is a recursive method, so the interpreter keeps a call stack. Each active call has its own parameters, node, nl, nr, l and r, stored in a frame on the stack. A call that asks its children pushes two new frames. When a child returns, its frame is removed and the parent continues with the next line.

The depth of the stack is the height of the tree, which is about log2(n). For one million items that is about 20 frames. The default recursion limit is 1000, so the search has plenty of room. The recursion limit matters only when a recursive function goes deep, which a segment tree does not do when it is built correctly.

The list takes most of the memory. A tree over n items uses 4 * n slots. A list of four million references takes about 32 megabytes, as the first lesson measured. Each int object adds memory as well, unless the value is one of Python's small cached integers. The tree is built once and then only read, so the query does not allocate new tree memory.

The build also uses the call stack, and it goes as deep as the tree is tall. It visits every node once, so its time is O(n) and its memory for the tree is O(n).

Common mistakes

Using the wrong neutral value. For a sum, the no-overlap case must return 0. If you return 1, every partly covered node adds an extra 1 for each outside child, and the answers are too big. For other operations the neutral value is different. Choose the value that changes nothing when combined.

Writing the no-overlap test the wrong way. The condition is r < nl or nr < l. Writing l > nr or r < nl is also correct, because it is the same test with the two parts swapped. The dangerous mistake is writing nl < l or nr > r and returning 0, which discards any node that is only partly inside the range. That drops items silently.

Forgetting the full-overlap test. Without it, the search still recurses until it reaches leaves, and it returns the right answer. It is just slower, and it looks like a performance bug. The test is short, so keep it.

Using / where you need //. In Python 3, / always gives a float. The midpoint (0 + 5) / 2 is 2.5, and using it as a list index fails:

mid = 2.5
Traceback (most recent call last):
  File "example.py", line 4, in <module>
    print(a[mid])
          ~^^^^^
TypeError: list indices must be integers or slices, not float

The last line names the problem. A list index must be a whole number. The fix is //, which is floor division and gives 2.

Passing the range in the wrong order. The query query(3, 1) is not a valid range, because the left end is after the right end. The code does not raise an error. It returns 0, because every node is either outside the range or fails the other tests. Check the order of the arguments before the call, or add a check at the start of query:

if l > r:
    raise ValueError(f"empty range: l={l} is greater than r={r}")

Mixing up the node numbers and the list positions. The node number is the slot in the tree list, while the list position is the index in the data. Keep them apart in your head and in your variable names. The build uses node for the slot and l and r for the data positions.

Reading input when there is none. A program that calls input() when the input is empty stops with this error:

EOFError: EOF when reading a line

This lesson's programs have their data built in, so they do not read input. Your own programs may read input, and then you must check that enough lines are present.

Debugging walkthrough: a query that is slightly wrong

Suppose a teammate says a query returns the wrong number for one range, but not for others. Here is a method to find the bug without guessing. Start with the simplest failing case. Suppose query(1, 3) gives 16 instead of 15. Run the traced program with that exact range, and read the output.

Look for the first line that does not match the case you expect. In a good trace, every line is right. Node 1 covers the whole list, so it must be partial. Node 2 covers [0..2], so it must be partial. Node 8 covers index 0, so it must be no overlap. If a line says "full overlap" for a node that is only partly inside, the test is wrong, and you have found the bug. If the lines are all right but the sum is wrong, the bug is in how the answers are added, so check the return line.

Once you find the line, fix it, then run the exhaustive check program. It runs every range for every size from 1 to 12. It only counts the passes, so if the passed count is smaller than the tested count, add a print inside the failing branch to see the first wrong range. That two-step method, a small trace and then a full check, finds almost every query bug.

Variations and edge cases

A query where l is greater than r is not a valid range. Your code should either refuse it or return the neutral value. The tree itself does not catch this, and it may return a plausible but wrong number. Add a check at the start of the public query method if your input can contain such pairs.

A range that is only one item is the simplest case. The search still works, but it follows one path from the root to a leaf. For a long list that path has about 20 partly covered nodes, and each of them makes a couple of calls that return at once. Remember this when you test. A single-item query is a good test of the leaves.

A range that covers the whole list returns the stored root sum at once. The first test in the search, the no-overlap test, is false, and the full-overlap test is true. So the root is returned without recursion. This is the fastest query possible.

A query can ask for a single answer that is not a sum. For example, a count of values above a threshold is a valid combine rule when each node stores a count. The shape of the code does not change. Only the stored value, the combine rule and the neutral value change. That is the main idea behind the generic version in the next lesson.

Interview questions this lesson prepares you for

"What is the complexity of building and querying?" The build is O(n) because every node is filled once. A query is O(log n) because at most two nodes per level are partly covered. Give the argument in two sentences, and interviewers will be satisfied.

"Why is the query not O(n)?" Because the search stops at any node that is fully inside or fully outside the range. Only the partly covered nodes recurse, and they form at most two paths down the tree. A candidate who explains the two paths is answering well.

"What does the stack depth look like?" It is about log2(n), because the recursion follows one level per call. Mention that this is why recursion is safe here, even for large lists.

"How would you test this?" Describe the exhaustive check: build trees of small sizes, test every range, and compare with a loop. Add a note that the loop is the reference, which is slow but obviously right.

Two more traces to practise on

Before you move on, trace two more queries on the same tree. Change the last line of the traced program to st.query(2, 4), and read the output line by line against the rules. The range is indexes 2 to 4, which holds the values 5, 7 and 9, so the answer should be 21. The output is:

node 1 [0..5]: partial overlap, ask both children
node 2 [0..2]: partial overlap, ask both children
node 4 [0..1]: no overlap, return 0
node 5 [2..2]: full overlap, return 5
node 3 [3..5]: partial overlap, ask both children
node 6 [3..4]: full overlap, return 16
node 7 [5..5]: no overlap, return 0
sum of 2..4 = 21

The root and node 2 are partly covered, so they ask their children. Node 4 covers [0..1], which lies entirely before index 2, so it returns 0 at once. Node 5 is the leaf for index 2, which is inside the range, so it returns 5. Node 3 is partly covered, and its child node 6 covers [3..4], which lies completely inside the range. Node 6 returns its stored sum, 16, and the search never enters its leaves. That is the saving the full-overlap case provides. Node 7 covers index 5, which is outside, and returns 0. The total is 5 + 16 + 0 = 21.

Count the calls. There are seven, which is the same size of work as the first trace, even though the range is different. The saving is visible in the missing calls. Had node 6 been treated as partly covered, the search would also have visited its two leaves, and it would have made nine calls instead of seven.

Now the whole list, which is the other edge of the range. Change the last line to st.query(0, 5). The output is:

node 1 [0..5]: full overlap, return 36
sum of 0..5 = 36

The whole list is a full overlap at the root, so the answer takes one call and no recursion at all. This is the fastest possible query, and it is a good check that the root holds the right total. If the root total is wrong, every query that touches the root is wrong too, and this one-line trace shows it immediately.

Recursion versus a loop

The query in this lesson is recursive, which means it calls itself. A loop could do the same job. Start at the root, keep a stack of nodes to check, and push the children that are partly covered. That version avoids the call stack, and it can be faster in very tight code.

The recursive version is easier to read, and it matches the tree picture. Each call is one node, and each case of the three-case rule is one if line. For learning, that match matters more than speed. When you can explain the recursive version, the iterative version is a small step.

A third option is the bottom-up iterative segment tree. It uses a tree size that is a power of two and walks up from the two ends of the query range. It is shorter and fast, and it is common in competitive programming. It is harder to extend to lazy updates, though, so learn the recursive version first. The lazy lesson depends on the recursive shape.

Both recursive and iterative versions run in O(log n) per query, so the choice does not change the complexity. It changes the constant factor and the readability. In an interview, write the recursive version unless the interviewer asks for the fastest possible code. Say that an iterative version exists, and you will show that you know the trade-off.

Explaining the bound to an interviewer, step by step

Interviewers often ask, "Why is the query fast?" A strong answer is a short argument with numbered steps. Here is one you can say out loud.

First, the search starts at the root and moves down one level at a time. There are about log2(n) levels in the tree, because each level halves the segment.

Second, a node can only recurse if it is partly covered. A partly covered node is one where the query range has an end inside the node's segment. Each end of the query range can cut through at most one node on each level, so at most two nodes per level are partly covered.

Third, a node that is fully inside or fully outside returns at once. It makes no further calls. So the only calls that lead to more work come from the at most two partly covered nodes per level.

Fourth, each partly covered node calls exactly two children. At most two partly covered nodes per level means at most four calls per level that do real work, and the other calls return at once.

Fifth, multiply. About log2(n) levels times about four calls per level gives about 4 * log2(n) calls. That is O(log n).

Practise saying this in about thirty seconds. The interviewer then knows you understand the reason for the speed, not only the result. You can add that the build is O(n) and every later query is O(log n), so the structure pays for itself after a few queries.

A real-world example: a sensor window

Suppose a weather station records the rainfall once per minute. A day has 1440 minutes. A dashboard asks, many times a day, for the total rainfall between minute 120 and minute 480. New readings arrive every minute, so the list keeps changing.

A loop would add up to 1440 numbers for each dashboard request. Prefix sums would need a rewrite of up to 1440 values when a reading is corrected. A segment tree handles both cheaply. Each query visits at most about 4 * log2(1440), which is about 42 calls, because log2(1440) is about 10.5. Each new reading is one point update that rewrites about 11 nodes on one path.

The same design works for any sensor that reports a value per time slot. Shops use it for hourly sales, and factories use it for the temperature of each machine. The key property is that the question is a range, the data changes, and the range is long enough that a loop is too slow. Whenever you see those three things together, think of a segment tree.

Practice

1. For n = 6, which single node is fully covered by query(0, 2)? Name the node number and its range.

Show the answer

Node 2, which covers [0..2]. The root [0..5] is partly covered, so it asks its children. Node 2 is inside [0..2] exactly, so it is full overlap and returns its stored sum. Node 3 covers [3..5], which is outside, so it returns 0 at once.

2. Build by hand the sum tree for [4, 1, 6]. Give the sum in each node and the node numbers.

Show the answer

The root is node 1, covering [0..2] with sum 11. Its left child is node 2, covering [0..1] with sum 5. Node 2 has leaves node 4 (value 4) and node 5 (value 1). The right child of the root is node 3, covering [2..2], which is a leaf with sum 6. The tree list needs slots 1 to 5 in use, so 4 * 3 = 12 slots is more than enough.

3. What neutral value should a maximum tree use for the no-overlap case? Give the Python value and say why.

Show the answer

Use float("-inf"), or any value smaller than every item you can store. A maximum combines by taking the larger value, and the neutral value must never win. Any real item is bigger than negative infinity, so a no-overlap node never changes the result.

4. For n = 6, how many calls does query(4, 4) make in total, counting calls that return at once? List the node numbers in call order.

Show the answer

Seven calls: nodes 1, 2, 3, 6, 12, 13 and 7. The root (node 1) is partly covered, so it calls node 2, which is outside and returns. It then calls node 3, which is partly covered. Node 3 calls node 6, which is partly covered. Node 6 calls node 12, which is outside, then node 13, which is inside. Node 3 finally calls node 7, which is outside.

5. Explain in your own words why the search must stop at a fully covered node instead of recursing into its children.

Show the answer

The node already holds the sum of every item in its segment, and every item is inside the query. Recursing into the children would add the same items again, through smaller nodes. The answer would be the same, but the work would grow. Stopping early is the reason the query is O(log n) and not O(n).

6. A teammate writes the no-overlap test as if nl < l or nr > r: return 0. Explain why this is wrong, using a partly covered node as the example.

Show the answer

Take the node [0..2] and the query [1..3]. The node is partly covered, because indexes 1 and 2 are inside. The test nl < l is 0 < 1, which is true, so the code returns 0 and throws away the items 1 and 2. A partly covered node must be split, not discarded. The correct test is r < nl or nr < l.

7. Write a method sum_from(l) that returns the sum from index l to the end of the list, using the query method. Assume the tree has an attribute n for the list length.

Show the answer

def sum_from(self, l):
    return self.query(l, self.n - 1)

The end of the list is index n - 1, so the range [l, n - 1] covers everything from l on. If l is 0, the result is the total of the list.

8. The exhaustive check tests sizes from 1 to 12. How many ranges does it test for size 12? Give the formula you used.

Show the answer

78. A list of size s has s * (s + 1) // 2 ranges, because each pair l less than or equal to r is one range. For s = 12 that is 12 * 13 // 2 = 78. The full program tests all sizes from 1 to 12, which is 364 ranges in total.

79. A tree is built over a single item [7]. What does _build do, and what does query(0, 0) return? Explain why this case is worth testing.

Show the answer

The build sees l == r at once, so it stores 7 in node 1 and does not recurse. The query query(0, 0) finds node 1 fully covered and returns 7. This case is worth testing because the code that splits a range never runs. A bug in the split logic cannot show up here, but a bug in the leaf or the list size can.

What you now know

• The build is recursive. It fills the leaves first and then each parent from its two children, which costs O(n).
• A query compares each node's segment with the range and takes one of three cases: no overlap returns 0, full overlap returns the stored sum, and partial overlap recurses into both children.
• Only partly covered nodes recurse, and there are at most two per level, so a query costs O(log n).
• The trace shows the exact order of the calls, and the counting program shows that a query over a million items visits about 75 nodes.
• Test the tree against a plain loop on every range of small sizes. Small cases catch off-by-one errors that a single example misses.

The next lesson changes a single item with a point update and adapts the same tree to answer minimum queries.