Problem Description

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.

Problem Link

Solution 1:

Binary Search Tree’s inorder traversal is a sorted array. We can use inorder traversal to find the two nodes that are swapped by mistake, and then swap them back.

Code

func recoverTree(root *TreeNode) {
    var prevNode, first, second *TreeNode
    inorderTraversal(root, &prevNode, &first, &second)
    if first != nil && second != nil {
        first.Val, second.Val = second.Val, first.Val
    }
}

func inorderTraversal(root *TreeNode, prevNode, first, second **TreeNode) {
    if root == nil {
        return
    }
    
    inorderTraversal(root.Left, prevNode, first, second)
    if *prevNode != nil && (*prevNode).Val > root.Val {
        if *first == nil {
            *first = *prevNode
        }
        *second = root
    }
    *prevNode = root
    inorderTraversal(root.Right, prevNode, first, second)
}