Binary Search Trees: Correct a BST where 2 nodes are swapped

Problem Statement:

You are given a root node of a binary search tree where 2 nodes are swapped.

You need to correct the BST and return it.

Solution

A simple solution is traverse the BST in in-order store the node values in an array and sort the array. Then insert the array into BST.

This will take O(nlogn) time.

We can solve it in O(n) time.

How do we do that?

We know there are 2 situations where the swapping might occur.

1. When we need to swap 2 adjacent nodes {1, 2, 4, 3, 5}

2. When we need to swap the nodes in different place like {5, 2, 3, 4, 1}

So we need to maintain the pointer for 4 locations:

1. The previous node (prev)

2. The current node (curr)

3. First number to be swapped (frist)

4. Second number to be swapped (second)

Now we start to traverse the tree. We will mark the first node as “prev” and move to the next node. Then we start making comparison.

When we see that the current node value is less than the previous node value, the prev node is the first target “first”.

Now we set the “second” pointer to the current node.

Then again we will start the loop and check the poorly placed node. If it is present , then we update the “second” pointer.

Else if we need to swap the 2 adj nodes, then “first” and “second” pointers are already in right place.

Solution in C++

#include<iostream>
#include<vector>
#include<string>
#include<unordered_set>

using namespace std;


struct Node
{
    int data;
    struct Node *left, *right;
};

Node *get_new_node(int item)
{
    Node *temp =  new Node;
    temp->data = item;
    temp->left = temp->right = NULL;
    return temp;
}

void print_inorder(Node* Node) 
{ 
    if (Node == NULL) 
        return; 
    print_inorder(Node->left); 
    cout<<" "<<Node->data; 
    print_inorder(Node->right); 
} 

void recover_tree_helper(Node* current, Node*& prev, Node*& first, Node*& second) 
{
  if (!current) return;
  recover_tree_helper(current->left, prev, first, second);
  if (prev) {
      if (current->data < prev->data) {
          if (!first) {
              first = prev;
          }
          second = current;
      }
  }
  prev = current;
  recover_tree_helper(current->right, prev, first, second);
}

void recover_tree(Node* root) 
{
  Node *prev = nullptr, *first = nullptr, *second = nullptr;
  recover_tree_helper(root, prev, first, second);
  swap(first->data, second->data);
  return;
}

int main()
{
    Node *root = get_new_node(6);
    root->left        = get_new_node(10);
    root->right       = get_new_node(2);
    root->left->left  = get_new_node(1);
    root->left->right = get_new_node(3);
    root->right->right = get_new_node(12);
    root->right->left = get_new_node(7);

    printf("Inorder Traversal before correction \n");
    print_inorder(root);
  
    recover_tree(root);
  
    printf("\nInorder Traversal after correction \n");
    print_inorder(root);

    return 0;
}

Output:

Inorder Traversal before correction
1 10 3 6 7 2 12
Inorder Traversal after correction
1 2 3 6 7 10 12