// This function does inorder traversal to find out the two swapped nodes.
// It sets three pointers, first, middle and last. If the swapped nodes are
// adjacent to each other, then first and middle contain the resultant nodes
// Else, first and last contain the resultant nodes
void correctBSTUtil( struct node* root, struct node** first,
struct node** middle, struct node** last,
struct node** prev )
{
if( root )
{
// Recur for the left subtree
correctBSTUtil( root->left, first, middle, last, prev );
// If this node is smaller than the previous node, it's violating
// the BST rule.
if (*prev && root->data < (*prev)->data)
{
// If this is first violation, mark these two nodes as
// 'first' and 'middle'
if ( !*first )
{
*first = *prev;
*middle = root;
}
// If this is second violation, mark this node as last
else
*last = root;
}
// Mark this node as previous
*prev = root;
// Recur for the right subtree
correctBSTUtil( root->right, first, middle, last, prev );
}
}
correctBSTUtil(struct BT *root, struct BT **f, struct BT **m, struct BT **l, struct BT **prev)
{
if(root)
{
correctBSTUtil(root->l, f, m ,l ,prev);
if(*prev && root->data< (*prev)->data)
{
if(!*f)
{
*f = *prev;
*m = root;
}else
{
*l = root;
}
}
*prev = root;
correctBSTUtil(root->r, f, m ,l ,prev);
}
}
void correctBSTUtil(struct node* root, struct node** first,struct node** middle, struct node** last,struct node** prev)
{
if (root)
{
correctBSTUtil(root->left, first, middle, last, prev);
if (*prev && root->data < (*prev)->data)
{
if (!*first)
{
*first = *prev;
*middle = root;
}
else
*last = root;
}
*prev = root;
correctBSTUtil(root->right, first, middle, last, prev);
}
}
void fix_bst(struct node *root)
{
struct node *first, *middle, *last, *prev;
first = middle = last = prev = NULL;
correctBSTUtil(root, &first, &middle, &last, &prev);
if (first && last)
swap(&(first->data), &(last->data));
else if (first && middle)
swap(&(first->data), &(middle->data));
}
void correctBST(struct BT *root)
{
struct BT *prev, *first, *last, *middle=NULL;
prev = first =last = middle=NULL;
correctBSTUtil(root, &first, &middle, &last , &prev);
if(last && first)
swap(&(first->data), &(last->data));
else if(first && middle)
swap(&(first->data), &(middle->data));
}
//--------------------------------------------------------------------------------------------------
// A function to fix a given BST where two nodes are swapped. This
// function uses correctBSTUtil() to find out two nodes and swaps the
// nodes to fix the BST
void correctBST( struct node* root )
{
// Initialize pointers needed for correctBSTUtil()
struct node *first, *middle, *last, *prev;
first = middle = last = prev = NULL;
// Set the poiters to find out two nodes
correctBSTUtil( root, &first, &middle, &last, &prev );
// Fix (or correct) the tree
if( first && last )
swap( &(first->data), &(last->data) );
else if( first && middle ) // Adjacent nodes swapped
swap( &(first->data), &(middle->data) );
// else nodes have not been swapped, passed tree is really BST.
}
