Recover Binary Search Tree

How to Solve Recover Binary Search Tree

Written By

Kenny Polyak

Kenny is a software engineer and technical leader with four years of professional experience spanning Amazon, Wayfair, and U.S. Digital Response. He has taught courses on Data Structures and Algorithms at Galvanize, helping over 30 students land new software engineering roles across the industry, and has personally received offers from Google, Square, and TikTok.

Jai Pandya

Jai is a software engineer and a technical leader. In his professional career spanning over a decade, he has worked at several startups and companies such as SlideShare and LinkedIn. He is also a founder of a saas product used by over 10K companies across the globe. He loves teaching and mentoring software engineers. His mentees have landed jobs at companies such as Google, Facebook, and LinkedIn.

## Recover Binary Search Tree Introduction

The Recover Binary Search Tree problem asks us to restore a binary search tree to its original form after two of its nodes have been swapped. This problem requires a strong understanding of the structure and properties of a binary search tree, namely that all nodes to the left of the root are smaller than the root and all nodes to the right are larger.

## Recover Binary Search Tree Problem

Example Inputs and Outputs

Example 1

Input:

    1
   /
  3
   \
    2

Output:

    3
   /
  1
   \
    2

Example 2

Input:

     4
    / \
   2   5
  / \   \
 6   3   1

Output:

     4
    / \
   2   5
  / \   \
 1   3   6

Constraints

## Recover Binary Search Tree Solutions

Overview and Intuition

In this question, we are provided with the root of a binary search tree, well, an almost binary search tree. Someone accidentally swapped two of its nodes, turning the BST into a regular binary tree. We are on a quest to find those two nodes. Once we find them, it shouldn't be difficult to swap them again to restore the original BST.

Python

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

In a BST, the value of any node is greater than the value of all nodes in its left subtree. In contrast, the value of any node is smaller than the value of all nodes in its right subtree. All nodes are related in this way to their respective left and right subtrees.

Now that we have some clarity about BST, let's return to our original problem. It says that two of the nodes of the given BST have been swapped. Let's look at what a BST might look like if that were the case.

Approach 1: Export to an Array - Two Pass Solution

Intuition

In the previous section, we learned about the idea of in-order traversal. We take it further and implement it in this approach. We traverse the given tree in an in-order fashion. When we process a node, we push its reference into a list. If this were a perfect BST, the list or array would be perfectly sorted, and all nodes would be in ascending order of their values. In our case, however, there are two defectors.

  1. When the defectors are next to each other in the resultant array.

Example: [1, 2, 4, 3, 5, 6]

We can see that 4 and 3, which are right next to each other, have swapped positions. The ascending order is broken only at one place, namely at the pair (4, 3), which are also the two nodes we are trying to find.

  1. When the defectors are at least one element apart in the output list.

Example: [6, 2, 3, 4, 5, 1]

In this case, the perfect sort is broken for two pairs of values - (6, 2) and (5, 1). It turns out that 5 and 1 are the errant nodes.

Algorithm

  1. We traverse the given tree in the in-order sequence, and push the processed nodes into a new list almostSorted. In the solution here, we use the recursive DFS algorithm to simplify the implementation. We could also use the iterative version.
  2. If there are N elements in the list, iterate the index i through almostSorted from 0 till N - 2.
  3. For each pair almostSorted[i] and almostSorted[i + 1], if they are not in ascending order, we note them in a variable swapped. If swapped already contains two nodes, it means that we have already encountered a pair out of order. This means that almostSorted[i + 1] must be the second defector.
  4. Now that we have identified the two swapped elements, we just need to swap their values again.
  5. Return the root of the binary tree, which has now become a perfect BST.

Code and Implementation

def recover_bst(self, root: Optional[TreeNode]) -> None:
    almost_sorted = []
    swapped = None

# traverse the given tree in in-order fashion
    # and populate `almostSorted` array
    def in_order(node: Optional[TreeNode]):
        if node is None:
            return

in_order(node.left)
        almost_sorted.append(node)
        in_order(node.right)

in_order(root)

for i in range(len(almost_sorted) - 1):
        if almost_sorted[i].val > almost_sorted[i + 1].val:
            if swapped is None:
                swapped = [almost_sorted[i], almost_sorted[i + 1]]
            else:
                swapped[1] = almost_sorted[i + 1]

swapped[0].val, swapped[1].val = (swapped[1].val, swapped[0].val)
    return root

Time/Space Complexity

Approach 2: In-order Recursive Traversal - Single Pass

Intuition

A tiny optimization over the naive approach can save us a pass through all nodes. We use a global variable lastProcessed to store the last processed node.

  1. Traverse the tree in order. We use a recursive DFS algorithm here. Each node compares its value to its in-order predecessor.
  2. Store the pair in the array swapped if the predecessor's value is greater than the current node.
  3. At the end, we swap the values of the nodes contained in swapped and return the root of the tree.

Code and Implementation

def recover_bst(self, root: Optional[TreeNode]) -> None:
    swapped = last_processed = None

def find_swapped_pair(current: Optional[TreeNode]):
        nonlocal swapped, last_processed
        if current is None:
            return
        find_swapped_pair(current.left)

if last_processed is not None:
            if last_processed.val > current.val:
                if swapped is None:
                    swapped = [last_processed, current]
                else:
                    swapped[1] = current

last_processed = current
        find_swapped_pair(current.right)

find_swapped_pair(root)
    swapped[0].val, swapped[1].val = swapped[1].val, swapped[0].val
    return root

Time/Space Complexity

Approach 3: In-order Iterative Traversal - Single Pass

Intuition

We have already seen a recursive implementation of the solution. We use an explicit stack in place of an implicit recursion stack to convert the solution to use iteration instead of recursion.

Code and Implementation

def recover_bst(self, root: Optional[TreeNode]) -> None:
    stack = []
    current = root
    last_processed = swapped = None

while stack or current:
        while current:
            stack.append(current)
            current = current.left

current = stack.pop()

if last_processed and last_processed.val > current.val:
            if swapped:
                swapped[1] = current
                break
            else:
                swapped = [last_processed, current]

last_processed = current
        current = current.right

swapped[0].val, swapped[1].val = swapped[1].val, swapped[0].val
    return root

Time/Space Complexity

Bonus - Morris Traversal

This approach may not be very relevant from the perspective of an interview. In Morris traversal, the predecessor makes a temporary connection to the next node. This allows us to traverse without using additional space for a stack, helping to maintain the connection orientations at the end of processing.