Range sum queries in a Python BST

Why it matters

A range query asks about all the values between a low limit and a high limit. A shop might want the total of all orders between 100 and 500 rupees. A game might want the total score of every player between rank 40 and rank 60. Both questions have two bounds, and the answer depends only on the values that fall inside them.

A binary search tree (BST) keeps its values in order, so it can answer this kind of question without looking at every value. The BST rule lets you skip whole branches of the tree. This lesson uses that rule. It assumes you can insert values and walk a tree, which the first lesson of this module, "BST basics, insert and search in Python", covers.

By the end of this lesson you will be able to write a range sum in two ways, explain why one of them skips work, and measure the difference by counting the nodes each method visits.

The question in plain words

A range is a low bound and a high bound. In this lesson both bounds are inclusive, which means a value equal to a bound counts. The range from 25 to 65 contains 25 and it contains 65.

A range sum is the total of all the values in the tree that fall inside the range. Take the values 25, 30, 35, 40, 45, 50, 60 and 65. They all fall inside the range from 25 to 65. Their total is 350, because 25 + 30 + 35 + 40 + 45 + 50 + 60 + 65 = 350.

The examples in this lesson use one tree with 14 nodes. It is built by inserting these values, in this order: 50, 30, 70, 20, 40, 60, 80, 10, 25, 35, 45, 65, 75, 85. Here is its shape. Each line is one node, and the indentation shows which node it hangs under.

50
├── 30
│   ├── 20
│   │   ├── 10
│   │   └── 25
│   └── 40
│       ├── 35
│       └── 45
└── 70
    ├── 60
    │   └── 65
    └── 80
        ├── 75
        └── 85

Read the tree from the top. The root is 50. Its left child is 30, and everything under 30 is smaller than 50. Its right child is 70, and everything under 70 is larger than 50. The same rule holds at every node.

A subtree is a node together with everything below it. The subtree of 30 holds 30, 20, 10, 25, 40, 35 and 45. Each child starts its own subtree.

Method 1: visit every node

The simplest method looks at every node in the tree. For each node, it asks whether the value is inside the range. If it is, the value is added to a running total. This method does not use the BST rule at all, so it would also work on a tree whose values are in any order.

It is easiest to write as a recursive function. Recursion means a function that calls itself on a smaller piece of the problem. Here the smaller piece is a subtree. The sum of a tree is the node's own value (if it is in range) plus the sum of its left subtree plus the sum of its right subtree. An empty tree, written None, has a sum of 0, and that is where the calls stop.

The counters below are used by both methods. visits counts how many nodes a method looks at. seen records the values in the order they were looked at, so you can read the walk. Each function that changes visits says global visits first, because a plain assignment inside a function would create a new local name instead.

visits = 0
seen = []
 
 
def reset():
    global visits
    visits = 0
    seen.clear()

Here is the simple method with the counter added.

def range_sum_all(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    here = node.val if lo <= node.val <= hi else 0
    return here + range_sum_all(node.left, lo, hi) + range_sum_all(node.right, lo, hi)

Read the last line first. It is the rule in one sentence: this node's contribution, plus the answer for the left subtree, plus the answer for the right subtree. The middle line is the range check. The form a if condition else b gives the value when it is between lo and hi, and 0 otherwise. Python integers have no fixed size, so the total never wraps around, as it can in languages with fixed-size integers. The "Common mistakes" section shows what that means in practice.

On the 14-node tree with the range from 25 to 65, this method visits every node. The sum is 350 and the visit count is 14. The 6 values outside the range (10, 20, 70, 75, 80 and 85) are still looked at, because the method has no way to know they are outside without checking them.

Method 2: use the BST rule to skip

Pruning means cutting off a branch of a search because the branch cannot give an answer. The BST rule makes this possible. Look at one node, called v, and think about its two subtrees.

If v is below the low bound, then everything in its left subtree is smaller than v. That means everything on the left is also below the low bound, so none of it can count. Skip the left subtree. Only the right subtree can hold matches. The node v itself does not count either.
If v is above the high bound, then everything in its right subtree is larger than v. That means everything on the right is also above the high bound. Skip the right subtree. Only the left subtree can hold matches.
If v is inside the range, it counts. Both subtrees may still hold matches, so visit both.

A plain example helps. Imagine a shelf of numbered boxes, sorted from left to right. You want the boxes numbered from 25 to 65. You pick up box 20 and see that it is below 25. Every box to the left of box 20 has a smaller number, so you do not need to open any of them. You only need to look to the right of box 20.

Here is the pruned method.

def range_sum_pruned(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    if node.val < lo:
        return range_sum_pruned(node.right, lo, hi)
    if node.val > hi:
        return range_sum_pruned(node.left, lo, hi)
    return node.val + range_sum_pruned(node.left, lo, hi) + range_sum_pruned(node.right, lo, hi)

The two if lines near the top are the pruning. The first one handles a node below the range and skips its left side. The second one handles a node above the range and skips its right side. The last line handles a node inside the range, which counts and visits both sides.

Trace the pruned method

Run the pruned method on the 14-node tree with the range from 25 to 65. The table shows each step.

Step	Node	Compared with 25 to 65	Action	Added
1	50	inside	add it, visit both sides	50
2	30	inside	add it, visit both sides	30
3	20	below 25	skip its left side (10), go right	0
4	25	inside	add it	25
5	40	inside	add it, visit both sides	40
6	35	inside	add it	35
7	45	inside	add it	45
8	70	above 65	skip its right side (75, 80, 85), go left	0
9	60	inside	add it, visit both sides	60
10	65	inside	add it	65

Step 3 is the first skip. The node 20 is below 25, so the value 10 under it is never looked at. Step 8 is the second skip. The node 70 is above 65, so the three values 75, 80 and 85 under its right side are never looked at.

The pruned method visited 10 nodes. The four skipped nodes are 10, 75, 80 and 85. The total is 25 + 30 + 35 + 40 + 45 + 50 + 60 + 65 = 350. Both methods give the same answer.

Here is the real output from running the program. The pruned visit order is the list of values in the order they were looked at.

small tree, range [25, 65]
  visit-all sum = 350, nodes visited = 14
  pruned sum    = 350, nodes visited = 10
  pruned visit order = [50, 30, 20, 25, 40, 35, 45, 70, 60, 65]

Why skipping is safe

Skipping only works because of the BST rule. Without the rule, the left subtree of a node could hold any value, and a skip could hide a match.

Here is a tree that breaks the rule. The root is 50. Its left child is 30, and 30 has a right child of 60. The value 60 is in the left subtree of 50, but 60 is larger than 50, so the tree is not a valid BST. Now run the pruned method for the range from 55 to 70. At the root, 50 is below 55, so the method skips the left subtree and goes right. The right side is empty. The method returns 0, which is wrong, because 60 is inside the range.

This is why the rule matters. The pruned method is correct for a valid BST, and it can give a wrong answer for a tree that is not one. If you build your tree with insert only, the rule always holds.

The output of that check is:

bad tree pruned [55,70]: 0

Counting visits: the real saving

The 14-node example is too small to show the difference. A larger test makes the saving clear. The program at the end of this lesson inserts 1000 random numbers between 0 and 9999 into a tree, using seed 42 so the result is the same every time. Duplicates are dropped, so the tree has fewer than 1000 nodes. The range is narrow: 4000 to 4100.

random tree built from 1000 values (duplicates are dropped), range [4000, 4100]
  visit-all sum = 64758, nodes visited = 961
  pruned sum    = 64758, nodes visited = 30
  brute force sum = 64758

The tree has 961 nodes, which is the visit count of the simple method. The pruned method visited 30 nodes and got the same sum. The third line is a brute force check: a plain sum over a list of values, with no tree at all. Brute force scans a whole list, so it is no faster than the simple method, but it is simple enough to trust, so it makes a good reference for checking a faster method.

Why 30? The pruned method visits two kinds of nodes. The first kind is the nodes whose values are inside the range. The second kind is the nodes on the way down toward the low bound and toward the high bound, because the walk has to pass through them to know which side to take. The range from 4000 to 4100 holds only a few of the tree's values, so most of the 30 visits are nodes on the two search paths. The count grows with the number of values in the range plus the height of the tree, not with the total number of nodes.

The height of a tree is the length of its longest path from the root to a leaf, counting nodes. A tree built from random values is usually short compared with its size, so the two paths are short. If the values were inserted in sorted order, the tree would become a long chain, and the two paths would be as long as the whole tree. The lesson "Why sorted input makes a Python BST a chain" covers that case, and this method is slow on it.

Worked example 1: an empty range

The first example is a range that holds no values. Use the range from 36 to 39 on the 14-node tree. No value in the tree is at least 36 and at most 39, so the answer should be 0. The question is how few nodes the pruned method needs to see to find that out.

Step	Node	Compared with 36 to 39	Action
1	50	above 39	skip its right side, go left
2	30	below 36	skip its left side, go right
3	40	above 39	skip its right side, go left
4	35	below 36	skip its left side, go right; there is no right child, so stop

The walk visited 4 nodes, and the sum is 0. The simple method would have visited all 14. The skipped subtrees here are the left side of 30 (20, 10 and 25), the right side of 50 (70, 60, 65, 80, 75 and 85), and the right side of 40 (45). Those are the 10 nodes the walk never looked at.

The output from the program for this range is:

[36,39] sum=0 visits=4 order=[50, 30, 40, 35] | visit-all sum=0 visits=14

Worked example 2: a range that covers the whole tree

The second example is the opposite case. Use the range from 0 to 100. Every node in the tree is inside it, so no skip ever happens.

[0,100] sum=690 visits=14 order=[50, 30, 20, 10, 25, 40, 35, 45, 70, 60, 65, 80, 75, 85] | visit-all sum=690 visits=14

The pruned method visits all 14 nodes, the same as the simple method. Pruning helps only when part of the tree is outside the range. When the range is wide, most of the tree has to be read anyway. The total 690 is the sum of all 14 values.

This example shows a useful rule. The cost of the pruned method depends on how much of the tree the range cuts away, so it is cheapest for narrow ranges.

Worked example 3: counting instead of summing

The same skipping rule works for any question that needs values inside the range. Here is a version that counts the values instead of adding them. Only the return value changes: each node that is inside the range returns 1, plus the counts from its two subtrees.

def count_in_range(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    if node.val < lo:
        return count_in_range(node.right, lo, hi)
    if node.val > hi:
        return count_in_range(node.left, lo, hi)
    return 1 + count_in_range(node.left, lo, hi) + count_in_range(node.right, lo, hi)

Run it on the range from 25 to 65 on the same tree. The output is:

count [25,65] sum=8 visits=10 order=[50, 30, 20, 25, 40, 35, 45, 70, 60, 65]

There are 8 values in the range: 25, 30, 35, 40, 45, 50, 60 and 65. The walk visits the same 10 nodes as the pruned sum did, because the walk is the same. Only the answer changes. This is the pattern to remember: a range question on a BST is a skipping walk with a different combining step.

The full program

The functions above are parts of one file. The run block below joins them with the tree code and prints the outputs shown in this lesson, except the TypeError in Mistake 2.

Python · editable

import random

visits = 0
seen = []

class Node:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

def insert(node, val):
    if node is None:
        return Node(val)
    if val < node.val:
        node.left = insert(node.left, val)
    elif val > node.val:
        node.right = insert(node.right, val)
    return node

def reset():
    global visits
    visits = 0
    seen.clear()

def range_sum_all(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    here = node.val if lo <= node.val <= hi else 0
    return here + range_sum_all(node.left, lo, hi) + range_sum_all(node.right, lo, hi)

def range_sum_pruned(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    if node.val < lo:
        return range_sum_pruned(node.right, lo, hi)
    if node.val > hi:
        return range_sum_pruned(node.left, lo, hi)
    return node.val + range_sum_pruned(node.left, lo, hi) + range_sum_pruned(node.right, lo, hi)

def count_in_range(node, lo, hi):
    global visits
    if node is None:
        return 0
    visits += 1
    seen.append(node.val)
    if node.val < lo:
        return count_in_range(node.right, lo, hi)
    if node.val > hi:
        return count_in_range(node.left, lo, hi)
    return 1 + count_in_range(node.left, lo, hi) + count_in_range(node.right, lo, hi)

def buggy_sum(node, lo, hi):
    if node is None:
        return 0
    if node.val < lo:
        return buggy_sum(node.left, lo, hi)
    if node.val > hi:
        return buggy_sum(node.right, lo, hi)
    return node.val + buggy_sum(node.left, lo, hi) + buggy_sum(node.right, lo, hi)

def build(values):
    root = None
    for v in values:
        root = insert(root, v)
    return root

def compare(label, root, lo, hi):
    reset()
    a = range_sum_all(root, lo, hi)
    a_visits = visits
    reset()
    p = range_sum_pruned(root, lo, hi)
    p_visits = visits
    order = list(seen)
    return label, a, a_visits, p, p_visits, order

small = build([50, 30, 70, 20, 40, 60, 80, 10, 25, 35, 45, 65, 75, 85])

print("small tree, range [25, 65]")
_, a, av, p, pv, order = compare("small", small, 25, 65)
print(f"  visit-all sum = {a}, nodes visited = {av}")
print(f"  pruned sum    = {p}, nodes visited = {pv}")
print(f"  pruned visit order = {order}")

bad = build([50, 30])
bad.left.right = Node(60)
print(f"bad tree pruned [55,70]: {range_sum_pruned(bad, 55, 70)}")

rnd = random.Random(42)
big_values = [rnd.randrange(10000) for _ in range(1000)]
big = build(big_values)
reference = sorted(set(big_values))
print("random tree built from 1000 values (duplicates are dropped), range [4000, 4100]")
_, a, av, p, pv, order = compare("big", big, 4000, 4100)
print(f"  visit-all sum = {a}, nodes visited = {av}")
print(f"  pruned sum    = {p}, nodes visited = {pv}")
print(f"  brute force sum = {sum(v for v in reference if 4000 <= v <= 4100)}")

for lo, hi in [(36, 39), (0, 100)]:
    _, a, av, p, pv, order = compare("", small, lo, hi)
    print(f"[{lo},{hi}] sum={p} visits={pv} order={order} | visit-all sum={a} visits={av}")

reset()
c = count_in_range(small, 25, 65)
print(f"count [25,65] sum={c} visits={visits} order={seen}")

print(f"buggy {buggy_sum(small, 25, 65)}")
print(2000000000 + 2000000000)

Common mistakes

Mistake 1: skipping the wrong side

The most common error is to swap the two skips. The code then goes the wrong way when a node is out of range. Here is the wrong version, with the two lines swapped.

# Wrong: the directions are swapped.
def buggy_sum(node, lo, hi):
    if node is None:
        return 0
    if node.val < lo:
        return buggy_sum(node.left, lo, hi)
    if node.val > hi:
        return buggy_sum(node.right, lo, hi)
    return node.val + buggy_sum(node.left, lo, hi) + buggy_sum(node.right, lo, hi)

Run it on the same tree and range, from 25 to 65:

buggy 200

The answer should be 350, and the program printed 200 with no error. The trace shows why. The walk adds 50 and 30. Its branch through 30 adds 40, 35 and 45, which are correct. At 20, which is below 25, the buggy code goes left to 10. Both 20 and 10 are below 25, so that branch ends and the value 25 is never added. At 70, which is above 65, the buggy code goes right to 80, then to 85. All of them are above 65, so the walk ends and 60 and 65 are never added. The total is 50 + 30 + 40 + 35 + 45 = 200.

A wrong answer with no error is harder to find than a crash. The safeguard is to keep a brute force check, like the one in the random test above, and to compare the two methods on many random trees.

Mistake 2: a branch that returns nothing

A function that does not return a value on one path gives None back to its caller. Adding None to a number then stops the program. Here is the range sum with a bare return on the branch for values below the range:

def range_sum_bare_return(node, lo, hi):
    if node is None:
        return 0
    if node.val < lo:
        return
    if node.val > hi:
        return range_sum_bare_return(node.left, lo, hi)
    return node.val + range_sum_bare_return(node.left, lo, hi) + range_sum_bare_return(node.right, lo, hi)

On the same tree, the call range_sum_bare_return(small, 25, 65) stops with this message:

TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'

The message says that an integer was added to None. The line in the traceback is the line that added the two results. The fix is return 0 on that branch, so that every path gives back a number.

Mistake 3: forgetting that a skipped node does not count

Some learners skip a subtree correctly but still add the node that caused the skip. For example, in the walk above, node 20 is below the range and must not be added to the total. The pruned code handles this with a separate line for each case. The node is added only in the last line, which runs only when the node is inside the range. If you combine the cases into one line that always adds node.val, values outside the range get added and the answer is wrong.

Mistake 4: a large sum is not a problem here, but a size limit is

Python integers have no fixed size limit, so a sum such as the one below is exact:

print(2000000000 + 2000000000)

The output is:

4000000000

Languages with fixed-size integers would wrap this sum around to a wrong number, so the same code written in another language needs a wider type for the total. When you write the same method in another language, check the largest value its integer type can hold.

Practice

Each exercise has its answer below it. Try the exercise on paper before you read the answer.

Exercise 1: trace the range from 30 to 60

Use the pruned method on the 14-node tree with the range from 30 to 60. List the nodes it visits, in order, and give the sum.

Show the answer

The visit order is 50, 30, 20, 25, 40, 35, 45, 70, 60, 65. The visit count is 10. The sum is 260, because 50 + 30 + 40 + 35 + 45 + 60 = 260. The values 20 and 25 are visited but are not added, because they are below 30. The node 65 is visited but is not added, because it is above 60.

Exercise 2: a range that covers everything

How many nodes does the pruned method visit on the 14-node tree for the range from 0 to 100? Explain why.

Show the answer

It visits all 14 nodes. Every node is inside the range, so no node causes a skip, and both subtrees of every node are visited. The sum is 690.

Exercise 3: an empty range

Give the sum and the visit count for the range from 36 to 39, and list the visited nodes in order.

Show the answer

The sum is 0. The visit count is 4. The visited nodes are 50, 30, 40 and 35. The walk goes left at 50, right at 30, left at 40, and then right at 35, where there is no child, so it stops.

Exercise 4: an empty range with low above high

The range from 60 to 20 has a low bound that is larger than the high bound. What does the pruned method return, and how many nodes does it visit?

Show the answer

It returns 0 and visits 3 nodes: 50, 70 and 60. At 50, the value is below 60, so the walk goes right to 70. At 70, the value is above 20, so the walk goes left to 60. At 60, the value is not below 60 and is above 20, so the walk goes left to nothing and stops. No value can be at least 60 and at most 20 at the same time, so the answer 0 is correct. The real output is [60,20] sum=0 visits=3 order=[50, 70, 60].

Exercise 5: change the sum into a count

Write a function that returns how many values in the tree are inside the range, using the skipping rule. Then state the answer for the range from 25 to 65 on the 14-node tree.

Show the answer

Use the count_in_range function from worked example 3. For the range from 25 to 65, it returns 8. The walk visits 10 nodes, in the order 50, 30, 20, 25, 40, 35, 45, 70, 60, 65.

Summary

A range sum adds every value between two inclusive bounds.
The simple method visits every node. It works on any binary tree and always costs one visit per node.
The BST rule lets you skip a whole subtree: a node below the low bound has a left subtree that is also below it, and a node above the high bound has a right subtree that is also above it.
The pruned method visits the nodes inside the range and the nodes on the two search paths toward the bounds. On a 961-node tree, a narrow range needed 30 visits instead of 961.
Check a fast method against a brute force sum, and make sure every path of a recursive function returns a number.

Range sum queries in a Python BST

Why it matters

The question in plain words

Method 1: visit every node

Method 2: use the BST rule to skip

Trace the pruned method

Why skipping is safe

Counting visits: the real saving

Worked example 1: an empty range

Worked example 2: a range that covers the whole tree

Worked example 3: counting instead of summing

The full program

Common mistakes

Mistake 1: skipping the wrong side

Mistake 2: a branch that returns nothing

Mistake 3: forgetting that a skipped node does not count

Mistake 4: a large sum is not a problem here, but a size limit is

Practice

Exercise 1: trace the range from 30 to 60

Exercise 2: a range that covers everything

Exercise 3: an empty range

Exercise 4: an empty range with low above high

Exercise 5: change the sum into a count

Summary

Stuck? Ask the tutor