Recover a BST with two swapped nodes in Python

Why this matters

A binary search tree, or BST, is a tree where each node holds a number, and the numbers follow one rule. Every value in a node's left subtree is smaller than the node's value, and every value in its right subtree is bigger. A subtree is a node together with everything below it. The rule lets a program find a value by going left or right at each node, instead of checking every node in the tree.

Sometimes two values in a BST end up in each other's places. A bug in the code that updates the tree can do this, and so can damage to a saved copy of the tree. The shape of the tree stays the same, so a quick look at the picture may not show anything wrong. But searches start to fail, because a search goes left when the value it wants is on the right.

This lesson shows how to find the two swapped nodes in one pass through the tree, and how to put their values back. The method uses an inorder walk, which visits every node in a fixed order. If you have not seen tree traversals before, the binary trees module before this one covers them. This lesson uses the walk and does not explain it again in full.

The idea: an inorder walk gives a sorted list

An inorder walk visits a node's left subtree first, then the node itself, then its right subtree. Here is a correct BST that holds the values 1 to 7:

        4
      /   \
    2       6
   / \     / \
  1   3   5   7

Walk this tree in inorder. Start at the root, 4. Visit its left subtree first, which is the subtree rooted at 2. Inside that, visit 1, then 2, then 3. Then visit 4 itself. Then visit the right subtree rooted at 6, which gives 5, then 6, then 7. The full list is 1, 2, 3, 4, 5, 6, 7.

This is not luck. In a correct BST, every value in a left subtree is smaller than the node, and the walk visits the left subtree before the node. Every value in a right subtree is bigger, and the walk visits the right subtree after the node. So an inorder walk of a correct BST always gives an increasing list.

That gives us a plan. If the list should go up every time, then any place where it goes down points to a problem. We call a place where a value is smaller than the value just before it an inversion. A correct tree has no inversions. Swapping two values creates one or two inversions, and the values in those inversions are the ones that are wrong.

We do not need the whole list to find them. We only need to compare each value with the one before it in the walk. That is a single pass, and it needs no extra list.

Worked example 1: two nodes that are far apart

Take the correct tree above and swap the values 2 and 6. The node that held 2 now holds 6, and the node that held 6 now holds 2. The shape does not change.

        4
      /   \
    6       2
   / \     / \
  1   3   5   7

Walk this tree in inorder. The list is 1, 6, 3, 4, 5, 2, 7. Compare each value with the one before it:

• 1 then 6. The list goes up. Fine.
• 6 then 3. The list goes down. This is the first inversion.
• 3 then 4, then 4 then 5. Both go up. Fine.
• 5 then 2. The list goes down. This is the second inversion.
• 2 then 7. The list goes up. Fine.

The first inversion is 6 and 3. The value 6 is the earlier one. It is wrong, because it sits where 2 belongs. The second inversion is 5 and 2. The value 2 is the later one, and it is wrong, because it sits where 6 belongs. The value 3 looked out of place only because 6 was in front of it.

So for two swapped values that are far apart, the rule is this. The first wrong value is the earlier value of the first inversion. The second wrong value is the later value of the last inversion. Here those are 6 and 2. Swap them back, and the list becomes 1, 2, 3, 4, 5, 6, 7 again.

Worked example 2: two nodes that sit next to each other

Now start from the correct tree and swap the values 3 and 4. The root now holds 3, and the node that used to hold 3 now holds 4:

        3
      /   \
    2       6
   / \     / \
  1   4   5   7

The inorder list is 1, 2, 4, 3, 5, 6, 7. Compare neighbours. The pairs 1 and 2, 2 and 4, and 3 and 5 all go up. The pair 4 and 3 goes down. That is one inversion. The earlier value, 4, is the first wrong value. The later value, 3, is the second wrong value. Swap them, and the list is correct again.

This case has only one inversion, so the first and the last inversion are the same one. That is why the rule has two parts. Set first when you see the first inversion, and set second at every inversion. Then the last value written into second is the right one.

Why write second again and again? Go back to worked example 1. Its first inversion gives 3 as a candidate for second, but 3 is correct, and the swap partner is 2. The later inversion, 5 then 2, replaces 3 with 2. The last write wins, and the last write is the right one. Here is the same walk with the values that each step would record:

inversion 1: 6 then 3   first = 6, second = 3
inversion 2: 5 then 2   second = 2  (replaces 3)

Worked example 3: the smallest and the largest values

Now swap the values 1 and 7. The 7 moves to the far left and the 1 moves to the far right:

        4
      /   \
    2       6
   / \     / \
  7   3   5   1

The inorder list is 7, 2, 3, 4, 5, 6, 1. There are two inversions. The first is 7 then 2, so first is 7 and second is 2. The second is 6 then 1, so second becomes 1. Swap the values of 7 and 1, and the list is 1 to 7 in order.

The two wrong nodes are at opposite ends of the tree and of the list. The method does not care how far apart they are. It only checks each value against the value next to it in the walk.

The code

Some words first. A class is a blueprint for a kind of object. Our TreeNode class holds one number in val, and two links called left and right. A link can be None, which means there is no node there.

A stack is a list where you add and remove items at the same end, like a pile of plates. In Python, a plain list works as a stack. append puts an item on top, and pop takes the top item off. The iterative version uses a stack to remember nodes that it has passed but not yet visited.

The program below has five parts. The leaf function builds a node with no children. The inorder function returns the list of values in inorder, so we can print the tree before and after the fix. The recover_bst function is the stack version of the lesson. The recursive version, walk and recover_recursive, does the same job with function calls. The main part builds three broken trees and runs the fix on each one, then runs the recursive fix on a fourth.

The real output is:

Example 1 before inorder: [1, 6, 3, 4, 5, 2, 7]
Example 1 after  inorder: [1, 2, 3, 4, 5, 6, 7]
Example 2 before inorder: [1, 2, 4, 3, 5, 6, 7]
Example 2 after  inorder: [1, 2, 3, 4, 5, 6, 7]
Example 3 before inorder: [7, 2, 3, 4, 5, 6, 1]
Example 3 after  inorder: [1, 2, 3, 4, 5, 6, 7]
Recursive after  inorder: [1, 2, 3, 4, 5, 6, 7]

Now read recover_bst one step at a time. The inner while loop goes left as far as it can, and pushes each node it passes onto the stack. When it cannot go further, stack.pop() takes the most recent node off the stack. That node is the next one in inorder. The check prev.val > cur.val is the inversion test from the idea section. The check prev is not None skips the test on the very first node, because the first node has nothing before it. The line if first is None: makes sure first is set only at the first inversion. The line second = cur runs at every inversion, so the last inversion wins. Then prev = cur moves the window forward, and cur = cur.right moves into the right subtree.

After the walk, the code swaps the two values only if it found both nodes. The line first.val, second.val = second.val, first.val swaps two values in one step. Python evaluates the right side first, so both values are read before either one is written. If the tree was correct to begin with, first and second stay None, and nothing changes.

The recursive version does the same steps in the same order. walk goes into the left subtree, checks the node against prev_r, and then goes into the right subtree. The three names first_r, second_r and prev_r are module-level variables. The global line tells Python that the function changes those shared names, so every call of walk sees the same values. The recursive version is easier to read, but each call adds a frame, the record Python keeps for that call, to the call stack. Python stops a call chain at about 1000 calls by default, so a very deep tree fails with a RecursionError. The stack version keeps its nodes in a list, which grows as needed, so it does not have that limit.

The program swaps values, not nodes. Each node keeps its place in the tree, and only the numbers move. Moving the nodes themselves would mean changing the links of their parents, and the values alone are enough to fix the tree.

A random check

Examples are good for understanding, but they only test the cases you thought of. A better check is to make many random trees, break each one, and see whether the fix restores it. With random.seed(42), 1000 random trials were made. Each trial inserted between 2 and 31 distinct values (random.randint(2, 31) picks the count, and random.sample(range(1, 200), n) picks the values) into a BST, then picked two nodes with random.choice. Trials where both picks were the same node (96 of them) were skipped. The other 904 swapped the values of the two nodes, ran recover_bst, and compared the inorder list with the sorted list of values. The check found 0 failures.

You can run this check yourself by adding a loop around the three examples. The one rule to keep is that the swapped pair must be two different nodes, so skip a trial when the two random positions are the same. The exact counts depend on the random numbers, so yours may differ from these.

Common mistakes

Forgetting the check for a missing previous node

The first node in the walk has no node before it, so prev is None at that point. If the code reads prev.val before it checks for None, Python stops the program. Here is the real error from a version with the prev is not None check removed, run on example 1:

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

The fix is to keep prev is not None and at the start of the test.

Setting first to the later value

The first wrong value is the earlier one of the first inversion. If the code writes first = cur instead of first = prev, it picks the later value, which is 3 in example 1. Here is the output of that version. The walk ends with 3 and 2 swapped, and the result is still wrong:

wrong first: [1, 6, 2, 4, 5, 3, 7]

Setting second only at the first inversion

If the code writes second only when first is still None, it keeps the value 3 from the first inversion and never sees the second inversion. Here is the output of that version. Example 1 ends with 6 and 3 swapped, and example 3 ends with 7 and 2 swapped:

wrong second ex1: [1, 3, 6, 4, 5, 2, 7]
wrong second ex3: [2, 7, 3, 4, 5, 6, 1]

The fix is the line second = cur outside the if first is None: block, as the code above has it.

Comparing nodes instead of values

prev and cur are nodes, not numbers. Python does not know how to compare two node objects with >, so it stops with this message:

TypeError: '>' not supported between instances of 'TreeNode' and 'TreeNode'

The fix is to compare the numbers inside them, with prev.val > cur.val.

Misspelling a field name

Python looks up attribute names when the code runs, so a typo such as first.value in place of first.val only fails on the line that reads it. The message is:

AttributeError: 'TreeNode' object has no attribute 'value'

The message names the attribute that does not exist and the type of the object. Check the spelling against the TreeNode class. A linter, a tool that reads code and reports likely mistakes, can find this kind of mistake before the program runs.

Practice

Exercise 1

The inorder list of a broken BST is 1, 5, 3, 4, 2, 6. Which two values were swapped?

Show the answer

The sorted list should be 1, 2, 3, 4, 5, 6. The first inversion is 5 then 3, so first is 5. The last inversion is 4 then 2, so second is 2. The swapped values are 5 and 2. After the swap the list is 1, 2, 3, 4, 5, 6.

Exercise 2

The inorder list of a broken BST is 1, 2, 4, 3, 5. Which two values were swapped?

Show the answer

There is one inversion, 4 then 3. The swapped values are 4 and 3. After the swap the list is 1, 2, 3, 4, 5.

Exercise 3

The inorder list of a tree that is already correct is 1, 2, 3, 4, 5. What do first and second hold after recover_bst runs?

Show the answer

Both are None. The list never goes down, so neither first nor second is ever set. The check first is not None and second is not None is then false, and no values change.

Exercise 4

The inorder list of a broken BST is 8, 2, 3, 4, 5, 6, 7, 1. Which two values were swapped, and which inversion sets each one?

Show the answer

The swapped values are 8 and 1. The first inversion is 8 then 2, which sets first to 8 and second to 2. The second inversion is 7 then 1, which sets second to 1. The final swap uses 8 and 1.

Exercise 5

Why does recover_bst set prev = cur for every node, and not only at an inversion?

Show the answer

Each test must compare a value with the value directly before it in the walk. If prev stayed behind after a node with no inversion, the next comparison would use an old value. For the list 1, 6, 3, 4, the comparison 6 then 3 is the inversion. If prev had stayed at 1 after the value 6, the code would compare 3 with 1 and never see the inversion.

Summary

• An inorder walk of a correct BST gives an increasing list, so any step down in that list is an inversion.
• For two swapped nodes that are far apart, the first wrong value is the earlier value of the first inversion, and the second is the later value of the last inversion.
• For two swapped nodes that are next to each other, there is only one inversion, and both values come from it.
• Set first only at the first inversion, and set second at every inversion, so the last inversion wins.
• Swap the two values back. The nodes stay where they are, and the shape of the tree does not change.