leetcode #114 in cpp

Given a binary tree, flatten it to a linked list in-place.

For example,
Given

         1
        / \
       2   5
      / \   \
     3   4   6

The flattened tree should look like:

   1
    \
     2
      \
       3
        \
         4
          \
           5
            \
             6

Solution:

give the root of a tree, we are going to flatten this tree. The result should only contains right subtree. 

The example in the question shows that the order in the final result should following the pre-order traversal. That is, we traverse the root, then traverse the left subtree, and then we traverse the right subtree. If we are going to make a single linked list, then we have to make sure that right subtree always comes after left subtree. 

Thus we need to:

1. make right subtree left subtree's rightmost node's right child. (ensure right subtree always comes after left subtree)

2. make the new left subtree as the right subtree. And make left subtree as empty.(ensure the final result only contains right subtree)

3. we continue to flatten(new right subtree).

Code:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    void flatten(TreeNode* root) {
        makeFlatten(root);
    }
    void makeFlatten(TreeNode *&root){
        if(!root) return;
        if(!root->right) {//if root right is empty, move root left to root right and delete root left
            root->right = root->left; 
            root->left = NULL;
            makeFlatten(root->right);//continue to flatten next node
            return; 
        }
        if(!root->left){//if root left is empty, we continue to faltten root right
            makeFlatten(root->right);
            return; 
        }
        //store root right as r. move root left to root right, then make r as the right child the right most point in the new right subtree.
        TreeNode *r = root->right; 
        //find the rightmost point in the left subtree
        TreeNode *cur = root->left;
        while(cur->right){
            cur = cur->right;
        }//put leftsubtree as right subtree
        root->right = root->left; 
        //delete left subtree
        root->left = NULL;
        //make r as right child of rightmost point in the new right subtree
        cur->right = r; 
        //continue to flatten right subtree.
        makeFlatten(root->right);
    }
};