[leetcode] 114 Flatten Binary Tree to Linked List

/*
http://leetcode.com/onlinejudge#question_114

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
*/
/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    typedef std::pair<TreeNode*, TreeNode*> nodepair_t;

    void connect(TreeNode * a, TreeNode * b)
    {
        if(a == b) { return; }
        if(a) { a->right = b; }
        if(b) { b->left = NULL; }
    }
    nodepair_t flatten_impl(TreeNode *root)
    {
        if(!root) { return make_pair(root, root); }

        nodepair_t L = flatten_impl(root->left);
        nodepair_t R = flatten_impl(root->right);
        nodepair_t ret = make_pair(root, root);

        if(L.first)
        {
            connect(root, L.first);
            ret.second = L.second;
        }
        if(R.first)
        {
            connect(ret.second, R.first);
            ret.second = R.second;
        }
        root->left = NULL;
        ret.second->right = NULL;
        return ret;
    }

    void flatten(TreeNode *root) {
        flatten_impl(root);
    }
};