Copy Binary Tree with Random Pointer

similar with the leetcode problem: copy LinkedList with random pointer

From geeksforgeeks

code: (test cases from geeksforgeeks and algorithm self-implemented)

#include <iostream>
using namespace std;
 
/* A binary tree node has data, pointer to left child, a pointer to right
   child and a pointer to random node*/
struct TreeNode
{
    int key;
    struct TreeNode* left, *right, *random;
};
 
/* Helper function that allocates a new Node with the
   given data and NULL left, right and random pointers. */
TreeNode* newNode(int key)
{
    TreeNode* temp = new TreeNode;
    temp->key = key;
    temp->random = temp->right = temp->left = NULL;
    return (temp);
}
 
/* Given a binary tree, print its Nodes in inorder*/
void printInorder(TreeNode* node)
{
    if (node == NULL)
        return;
 
    /* First recur on left sutree */
    printInorder(node->left);
 
    /* then print data of Node and its random */
    cout << "[" << node->key << " ";
    if (node->random == NULL)
        cout << "NULL], ";
    else
        cout << node->random->key << "], ";
 
    /* now recur on right subtree */
    printInorder(node->right);
}
 
TreeNode *copy(TreeNode *node)
{
    if(node == NULL) return NULL;
    TreeNode *newnode = newNode(node->key);
    
    TreeNode *lc = node->left;
    node->left = newnode;
    newnode->left = lc;
    
    TreeNode *newleft = copy(lc);
    TreeNode *newright = copy(node->right);
    newnode->right = newright;
    return newnode;
}

void linkRandom(TreeNode *root)
{
    if(root == NULL) return;
    if(root->random)
        root->left->random = root->random->left;
    linkRandom(root->left->left);
    linkRandom(root->right);
}

TreeNode *separate(TreeNode *root)
{
    if(root == NULL) return NULL;
    TreeNode *newroot = root->left;
    root->left = newroot->left;
    
    newroot->left = separate(root->left);
    newroot->right = separate(root->right);
    
    return newroot;
}
 
/* Driver program to test above functions*/
int main()
{
/*  //Test No 1
    Node *tree = newNode(1);
    tree->left = newNode(2);
    tree->right = newNode(3);
    tree->left->left = newNode(4);
    tree->left->right = newNode(5);
    tree->random = tree->left->right;
    tree->left->left->random = tree;
    tree->left->right->random = tree->right;
 
//  Test No 2
//  Node *tree = NULL;
/*
//  Test No 3
    Node *tree = newNode(1);
 
//  Test No 4
    Node *tree = newNode(1);
    tree->left = newNode(2);
    tree->right = newNode(3);
    tree->random = tree->right;
    tree->left->random = tree;
 
  Test No 5
    Node *tree = newNode(1);
    tree->left = newNode(2);
    tree->right = newNode(3);
    tree->left->left = newNode(4);
    tree->left->right = newNode(5);
    tree->right->left = newNode(6);
    tree->right->right = newNode(7);
    tree->random = tree->left;
*/
//  Test No 6
    TreeNode *tree = newNode(10);
    TreeNode *n2 = newNode(6);
    TreeNode *n3 = newNode(12);
    TreeNode *n4 = newNode(5);
    TreeNode *n5 = newNode(8);
    TreeNode *n6 = newNode(11);
    TreeNode *n7 = newNode(13);
    TreeNode *n8 = newNode(7);
    TreeNode *n9 = newNode(9);
    tree->left = n2;
    tree->right = n3;
    tree->random = n2;
    n2->left = n4;
    n2->right = n5;
    n2->random = n8;
    n3->left = n6;
    n3->right = n7;
    n3->random = n5;
    n4->random = n9;
    n5->left = n8;
    n5->right = n9;
    n5->random = tree;
    n6->random = n9;
    n9->random = n8;
 
/*  Test No 7
    Node *tree = newNode(1);
    tree->left = newNode(2);
    tree->right = newNode(3);
    tree->left->random = tree;
    tree->right->random = tree->left;
*/
    cout << "Inorder traversal of original binary tree is: \n";
    printInorder(tree);
 
    TreeNode *clone = copy(tree);
    linkRandom(tree);
    clone = separate(tree);
 
    cout << "\n\nInorder traversal of cloned binary tree is: \n";
    printInorder(clone);
 
    return 0;
}