Trim a BST to a range in Python

Why this matters

Many programs only need the values between two limits. A shop might show only the products priced between 10 and 50 dollars. A calendar might show only the events in one week. A search might keep only the results with a date between two days. The tree is a binary search tree (BST), so we can throw away the parts that are out of range and keep a smaller BST that holds only the values in range.

This operation is called trimming. It changes the tree itself, so the removed nodes are gone afterwards. The lesson on range sum queries in a BST asks a similar question, the total of the values in a range, but it does not change the tree. Trimming is the right tool when the same range will be used again and again, and the tree should stay small.

A range here is two numbers, low and high. A value is in range when low is less than or equal to the value, and the value is less than or equal to high. Both ends count. We assume low is not bigger than high.

The idea: three kinds of node

Look at any node in the tree. Its value falls into one of three groups.

Too small: the value is less than low. Every value in the left subtree is even smaller, so the node and its whole left subtree must go. The right subtree may still hold values in range, so it has to be kept.
Too big: the value is greater than high. Every value in the right subtree is even bigger, so the node and its whole right subtree must go. The left subtree may hold values in range, so it is kept.
In range: the node stays. Its left subtree and its right subtree each need trimming on their own.

The third rule is the one that makes the method work. Each subtree is trimmed the same way, and the result of trimming a child is attached back to its parent.

Here is a tree that holds 1, 3, 4, 5, 7, 8 and 9. Trim it to the range 4 to 8:

        5
      /   \
    3       8
   / \     / \
  1   4   7   9

Start at the root, 5. It is in range, so it stays. Its left child is 3, which is too small. Node 3 goes, and so does its left child 1, because 1 is smaller still. Node 3's right child is 4. It is in range, so it stays and takes 3's place. The root's right child is 8. It is in range, so it stays. Its left child 7 is in range, so it stays. Its right child 9 is too big, so it goes.

The trimmed tree is:

Its inorder walk is 4, 5, 7, 8. Inorder means left subtree, then node, then right subtree, and it gives the values in increasing order. The four values are exactly the values in range, and the tree still follows the BST rule.

The recursive function

Recursion means a function that calls itself on smaller pieces of the problem. Here the piece is a subtree. The function takes a node and returns the root of the trimmed subtree. The caller then stores that returned node as its new child. That is how a change deep in the tree reaches the parent.

The function follows the three groups from the previous section:

If the node is None, there is nothing to trim, so return None.
If the node is too small, return the trimmed right subtree. The node and its left subtree are dropped by not returning them.
If the node is too big, return the trimmed left subtree.
Otherwise the node is in range. Trim its left child and store the result as the new left child. Trim its right child and store the result as the new right child. Return the node.

Step 2 is the one that beginners get wrong. The node is removed, but its right child is still returned, so its right subtree is not lost.

Here is a worked run on the tree above, with range 4 to 8. The call starts at node 5. It is in range, so it calls itself on the left child, 3. Node 3 is too small, so the call returns the result of trimming node 4. Node 4 is in range, and its children are empty, so it returns itself. The root stores 4 as its new left child. The root then trims its right child, 8. Node 8 is in range, so it trims 7, which returns 7, and trims 9, which is too big and returns None. Node 8 stores 7 as its left child and None as its right child. The root returns itself. The final tree matches the picture above.

The iterative version

The recursive function is short, but it uses the call stack for every level of the tree. The iterative version uses a list as its own stack, so it works the same way with a loop. It has two steps.

The first step handles the root. If the root is out of range, the whole root must go, and the new root is one of its children. The rule is simple. If the root is too small, move to its right child. If the root is too big, move to its left child. Repeat until the node is in range or the node is None.

The second step walks through every node that is still in the tree and fixes its two child links. A child that is too small is replaced by its right child, and the loop checks again, because the new child might also be too small. A child that is too big is replaced by its left child, and the loop checks again for the same reason. Each node that is kept goes on the stack so its own children are checked.

Here is a tree where the root is out of range. Trim it to the range 5 to 9:

The root, 2, is too small. Move to its right child, 12. That is too big. Move to its left child, 6. Node 6 is in range, so the first step is done. Node 6 has a left child 4, which is too small, so the left link is replaced by 4's right child, which is None. Node 6 keeps 8 on the right. The trimmed tree holds 6 and 8, and its inorder walk is 6, 8.

The two loops in the second step use while, not if. The reason is that replacing a child can give a new child that is also out of range. The common mistakes section shows what happens when the loop is an if.

Here is the core of the iterative version, so you can see the two steps:

# Step 1: move the root until it is in range.
while root is not None and (root.val < low or root.val > high):
    root = root.right if root.val < low else root.left
 
if root is None:
    return None
 
# Step 2: fix the child links of every node that is still in the tree.
stack = [root]
while stack:
    node = stack.pop()
    while node.left is not None and node.left.val < low:
        node.left = node.left.right
    while node.right is not None and node.right.val > high:
        node.right = node.right.left
    if node.left is not None:
        stack.append(node.left)
    if node.right is not None:
        stack.append(node.right)

A complete program with three examples

The program below builds three trees and trims each one. The first example uses the recursive function, and the second and third use the iterative version, so both are run on real input. The last part of the program checks both versions against a simple filter on 2000 random trees.

The leaf function builds a node with no children. The inorder function returns the values in increasing order, so the output shows the result of each trim. The insert function is an ordinary BST insert, used to build the random trees.

Python · editable

import random

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

def leaf(val):
    return TreeNode(val)

# Recursive trim: returns the new root of the subtree.
def trim_recursive(node, low, high):
    if node is None:
        return None
    if node.val < low:
        # Node and its left subtree are all too small. Keep only the right subtree.
        return trim_recursive(node.right, low, high)
    if node.val > high:
        # Node and its right subtree are all too large. Keep only the left subtree.
        return trim_recursive(node.left, low, high)
    node.left = trim_recursive(node.left, low, high)
    node.right = trim_recursive(node.right, low, high)
    return node

# Iterative trim: first move the root into range, then fix the child links.
def trim_iterative(root, low, high):
    while root is not None and (root.val < low or root.val > high):
        root = root.right if root.val < low else root.left
    if root is None:
        return None
    stack = [root]
    while stack:
        node = stack.pop()
        while node.left is not None and node.left.val < low:
            node.left = node.left.right
        while node.right is not None and node.right.val > high:
            node.right = node.right.left
        if node.left is not None:
            stack.append(node.left)
        if node.right is not None:
            stack.append(node.right)
    return root

def inorder(root):
    out = []
    stack = []
    cur = root
    while cur is not None or stack:
        while cur is not None:
            stack.append(cur)
            cur = cur.left
        cur = stack.pop()
        out.append(cur.val)
        cur = cur.right
    return out

def insert(root, v):
    if root is None:
        return leaf(v)
    if v < root.val:
        root.left = insert(root.left, v)
    else:
        root.right = insert(root.right, v)
    return root

def build_example1():
    return TreeNode(5, TreeNode(3, leaf(1), leaf(4)), TreeNode(8, leaf(7), leaf(9)))

def build_example2():
    return TreeNode(2, leaf(1), TreeNode(12, TreeNode(6, leaf(4), leaf(8)), leaf(15)))

def build_example3():
    return TreeNode(6, TreeNode(2, leaf(1), TreeNode(4, leaf(3), leaf(5))), TreeNode(8, None, None))

# Example 1: range [4, 8]. Node 3 is too small, but its right child 4 stays.
print("Example 1 before inorder:", inorder(build_example1()))
print("Example 1 trimmed inorder:", inorder(trim_recursive(build_example1(), 4, 8)))

# Example 2: range [5, 9]. The root 2 is too small, so the root moves down.
print("Example 2 before inorder:", inorder(build_example2()))
print("Example 2 trimmed inorder:", inorder(trim_iterative(build_example2(), 5, 9)))

# Example 3: range [3, 7]. Node 2 is too small, so it is replaced by its right child 4.
print("Example 3 before inorder:", inorder(build_example3()))
print("Example 3 trimmed inorder:", inorder(trim_iterative(build_example3(), 3, 7)))

# Random check against a plain list filter, for both versions.
rng = random.Random(7)
failures = 0
for _ in range(2000):
    n = 1 + rng.randrange(40)
    values = list(range(1, n + 1))
    rng.shuffle(values)
    low = 1 + rng.randrange(n)
    high = low + rng.randrange(n - low + 1)
    want = list(range(low, high + 1))
    a = None
    b = None
    for v in values:
        a = insert(a, v)
        b = insert(b, v)
    ra = trim_recursive(a, low, high)
    rb = trim_iterative(b, low, high)
    if inorder(ra) != want or inorder(rb) != want:
        failures += 1
print("random failures:", failures)

Save the code as trim.py, then run python3 trim.py. The real output is:

Example 1 before inorder: [1, 3, 4, 5, 7, 8, 9]
Example 1 trimmed inorder: [4, 5, 7, 8]
Example 2 before inorder: [1, 2, 4, 6, 8, 12, 15]
Example 2 trimmed inorder: [6, 8]
Example 3 before inorder: [1, 2, 3, 4, 5, 6, 8]
Example 3 trimmed inorder: [3, 4, 5, 6]
random failures: 0

Example 1 shows the recursive function on the tree from the idea section. The values 1, 3 and 9 are gone, and 4 stays in place of 3. The walk prints 4, 5, 7, 8, which is the range, in order.

Example 2 shows the root moving. The root 2 is too small, so the iterative version moves to 12, then to 6. The values 1, 2, 4, 12 and 15 are removed. The walk prints 6 and 8.

Example 3 shows a too-small node with a right child that stays. The node 2 is below the range 3 to 7, so it is replaced by its right child, 4. The value 1 goes with it, because 1 is in the left part of 2, and the value 4 stays. The node 4 keeps its own children, 3 and 5, because both are in range. The node 8 is too big, so it is removed. The final walk is 3, 4, 5, 6.

The random check builds 2000 trees, each with 1 to 40 distinct values. Each tree is trimmed to a random range, and both versions are compared with the plain list of values in that range. The result is 0 failures.

Common mistakes

Dropping the whole subtree of a node that is too small

A beginner may write return None for a node that is too small, because the node goes and its subtree seems to go with it. The right subtree can hold values in range, so this throws them away. Here is the output of a version that makes this change in the recursive function only. Example 1 loses the value 4:

Example 1 trimmed inorder: [5, 7, 8]
random failures: 1504

The random check found 1504 failures out of 2000 trees. The fix is the line return trim_recursive(node.right, low, high), which keeps the right subtree.

Using if instead of while in the iterative version

A replacement child can also be out of range. For example, the root 10 has a left child 2, 2's right child is 3, and 10's right child is 15. The range is 4 to 20, so both 2 and 3 are too small. If the code uses if to check the left child once, it replaces 2 with 3, and 3 is still too small. Here is the output of a version with if in place of both while loops in the second step:

before [2, 3, 10, 15]
if-bug trimmed [3, 10, 15]
correct trimmed [10, 15]

The value 3 should have been removed, and the version with if keeps it. On the random check, the same version failed 904 times out of 2000 trees. The fix is while, so the check runs again on the new child until it is in range or None.

Forgetting the empty tree

If every node is out of range, the first step of the iterative version moves the root to None. The stack step then tries to read a field of None, and Python stops the program. Here is the real error, from a version without the check if root is None: return None, run on a single node 1 with the range 5 to 9:

AttributeError: 'NoneType' object has no attribute 'left'

The fix is the check after the first loop. The function must return None when nothing is left.

Trimming a tree that is not a BST

Trimming relies on the BST rule. If the input breaks that rule, the output can break it too. A tree with root 5 and left child 6 is not a BST, because 6 is bigger than 5 but sits on its left. Trim it to the range 3 to 7. Both values are in range, so both stay, and the inorder walk comes out in the wrong order:

not a BST trimmed [6, 5]

A list that is not in increasing order is a sign that the input tree was wrong, not that the trim failed. Check the input with an inorder walk before you trim it.

Practice

Exercise 1

Trim the tree 7 (left 3 with children 1 and 5, right 9 with children 8 and 10) to the range 4 to 8. Which values are kept?

Show the answer

The values 5, 7 and 8 are kept. The root 7 is in range. Its left child 3 is too small, so it is replaced by its right child 5, which is in range. Its right child 9 is too big, so it is replaced by its left child 8, which is in range. The values 1, 3, 9 and 10 are removed. The trimmed tree is 7 with left child 5 and right child 8.

Exercise 2

What does trim_iterative return when every node is out of range?

Show the answer

It returns None. The first loop moves the root until it becomes None, and the check if root is None: return None returns before the stack is used.

Exercise 3

The tree has the root 6, its left child is 1 with no children, and its right child is 8 with left child 7. Trim it to the range 5 to 10. What is the inorder walk of the result?

Show the answer

The walk is 6, 7, 8. The root 6 is in range and stays. The left child 1 is too small and has no right child, so it becomes None. The right child 8 is in range and stays. Its left child 7 is in range and stays.

Exercise 4

Why does the code keep the left subtree of a node that is too big?

Show the answer

Every value in the left subtree is smaller than the node. The node is above high, but a value in the left subtree may still be at or below high and at or above low. For example, under the node 9 with the range 4 to 7, the left child 6 is in range, so it must stay.

Exercise 5

Trimming only removes nodes and changes child links. Why does the result still follow the BST rule without any rebalancing or sorting?

Show the answer

Removing a node never moves a value to the other side of a node. Each kept node has the same left and right groups as before, and those groups only lost values. A subtree that was all smaller than a node is still all smaller after some of its values are removed. So the BST rule still holds for the nodes that remain.

Summary

A node that is too small goes, but its right subtree may hold values in range, so keep it.
A node that is too big goes, but its left subtree may hold values in range, so keep it.
The recursive function returns the new subtree and stores it as the parent's child.
The iterative version first moves the root into range, then fixes each child link with a while loop.
Check the empty tree, because trimming can remove every node.

Trim a BST to a range in Python

Why this matters

The idea: three kinds of node

The recursive function

The iterative version

A complete program with three examples

Common mistakes

Dropping the whole subtree of a node that is too small

Using if instead of while in the iterative version

Forgetting the empty tree

Trimming a tree that is not a BST

Practice

Exercise 1

Exercise 2

Exercise 3

Exercise 4

Exercise 5

Summary

Stuck? Ask the tutor