Recover Binary Search Tree

You are given the root of a binary search tree (BST), where the values of exactly two nodes of the tree were swapped by mistake. Recover the tree without changing its structure. Follow up: A solution using O(n) space is pretty straight-forward. Could you devise a constant O(1) space solution?

Basic strategy

If we perform an inorder traversal, we can easily check where "broken" nodes are, by checking nodes against their neighbours. If we keep track of the bad nodes we find, the two that should be swapped will always be the first bad node and the last bad node we find.

The real challenge

Usually to do an inorder traversal, we'd either use recursion or a stack to visit all of the nodes. Both of those solutions require O(n) space though!

This was actually a team effort. One of the other members of my data structres and algorithms club found the "Morris Traversal", and shared this article with an implementation in Java.

Morris Traversal

This is a really cool algorithm, using a "threaded" binary tree. The tree is not threaded when you get it, but you add threads as you go. The basic approach is:

1. Check if the node has a left child. If it doesn't, visit it and go right. If there's no right child, you're done.
2. If there's a left child, get the "previous" node, meaning the one that comes before it in an inorder traversal. This is the rightmost child of the left child. "Thread" the tree by pointing the previous node to the node you're processing. Then go left.

So all you're doing is finding empty right nodes and changing thir pointers!

The getPrev function

This function gets the "previous" node defined above, but it makes sure that it doesn't return the node that was passed in.

Then after you call getPrev you can check if the returned node has a right node. If it does, that means the node you're processing has already been threaded. It's time to break the thread, process that node, and go right.

Solution

/**
 * @param {TreeNode} root
 * @return {TreeNode} The fixed tree.
 */
export default function recoverTree(root) {
    // Use the Morris Traversal to find bad nodes in O(n) time and 0(1) space
    // Morris Traversal uses a "Threaded Tree", where you modify the tree as you traverse it to save space.
    const inorder =[];
    const badNodes = []
    let cur = root
    while (cur) {
        if (cur.left === null) {
            visit(cur, inorder, badNodes)
            cur = cur.right
            continue
        }
        const prev = getPrev(cur)
        if (prev.right === null) {
            // This means it has not been threaded already
            prev.right = cur
            cur = cur.left
            continue
        }
        // We've hit a threaded node. Break the thread and go right
        prev.right = null
        visit(cur, inorder, badNodes)
        cur = cur.right
    }
    // Swap the two bad nodes.
    const firstBad = badNodes[0].val
    badNodes[0].val = badNodes[1].val
    badNodes[1].val = firstBad
    return root
}

/**
* Should only be called if node.left exists
* Returns the rightmost node of the left child, but not itself if it finds itself
*/
function getPrev(node) {
    // only call if `node.left` exists
    let cur = node.left
    while (cur.right && cur.right !== node) {
        cur = cur.right
    }
    return cur
}

/**
* Update inorder traversal and badNodes based on the next node.
*/
function visit(node, inorder, badNodes) {
    if (inorder.length === 0) {
        inorder.push(node)
        return
    } else if (inorder.length === 1) {
        inorder.push(node)
    } else {
        inorder[0] = inorder[1]
        inorder[1] = node
    }
    if (inorder[0].val > inorder[1].val) {
        // bad node found!
        if (badNodes.length === 0) {
            badNodes[0] = inorder[0]
        }
        badNodes[1] = inorder[1]
    }
}

Big O Analysis

Time Complexity: O(n)

Explanation: We should only have to scan each node once to find the out-of-place nodes.