Inverting a binary tree

Question

I have written this code to invert a binary tree.

Thus if the original tree is:

Then the tree returned should be:

/**
 * 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:
    TreeNode* invertTree(TreeNode* root) {
        if(root==NULL)
            return NULL;

        /* Performing post-order traversal of the tree */

        //Visiting the left node first
        if(root->left != NULL)
            invertTree(root->left);

        //Visiting the right node next
        if(root->right != NULL)
            invertTree(root->right);

        //Visiting the current node; swapping the left and the right nodes of the current node
        TreeNode* tn = new TreeNode(0);
        tn=root->right;
        root->right=root->left;
        root->left=tn;

        return root;
    }
};

While the code runs correctly and returns the required output, I think I am invoking undefined behavior when the control hits the leaf nodes. Since, at the leaf, a new node is created and it is assigned (incorrectly) the right of the current node and so on.

Could someone please clarify if what I am doing is indeed wrong?

Note: The question and the images are taken from LeetCode.com

Answer 1

Here you are creating a node unnecessarily:

   TreeNode* tn = new TreeNode(0);
    tn=root->right;

You could write simpler as:

   TreeNode* tn = root->right;

Because you're not using the original value of tn anyway, you only need it as a work variable for swapping the left and right nodes.

Also, your implementation doesn't use the return value of the recursive invertTree calls. You can use them better and get rid of the conditional statements, making the solution more compact:

TreeNode* invertTree(TreeNode* root) {
    if (root == NULL) {
        return NULL;
    }

    TreeNode* work = root->right;
    root->right = invertTree(root->left);
    root->left = invertTree(work);
    return root;
}

 I think I am invoking undefined behavior when the control hits the leaf nodes.

No, you're not invoking "undefined behavior" when hitting the leaf nodes. Actually I think you meant to say intermediary nodes. At leaf nodes, the function receives a null node and you return null. At intermediary nodes, you allocate an unnecessary node, as explained earlier. The unnecessary allocation doesn't trigger an undefined behavior, but since you never free the allocated memory, this is a memory leak, and indeed a problem if the program was used for a substantial amount of time, accumulating unreleased memory.

Answer 2

Apart from what the other answer said you can also use std::swap from STL.

Also you just need the first if statement, because if root->left or root->right is NULL it will return after one more recursive call anyway.

    void invertTree(TreeNode* root) {
        if(root == NULL) {
            return;
        }
        std::swap(root->left, root->right);
        invertTree(root->left);   
        invertTree(root->right);
    }

Answer 3

The shortes possible code:

    void invertTree(TreeNode* root) {
        if(root) {
            std::swap(root->left, root->right);
            invertTree(root->left);   
            invertTree(root->right);
        }
    }