LeetCode: Construct Binary Tree from Preorder and Inorder Traversal

Given preorder and inorder traversal of a tree, construct the binary tree.

Note:
You may assume that duplicates do not exist in the tree.

---

We can construct the tree recursively. The first node of preorder list will be the root node of the tree. Then we can find the root node in inorder list. And elements on the left side of the root node will construct the left subtree. Elements on the right side of the root node will construct the right subtree.

C++ Code:

/*
 * func: build_tree_preorder_inorder_helper
 * goal: helper function to build tree
 * @param preorder: preorder traversal list
 * @param pre_start: starting index
 * @param pre_end: ending index
 * @param inorder: inorder traversal list
 * @param in_start: starting index for inorder
 * @param in_end: ending index for inorder
 * return: root node
 */
TreeNode *build_tree_preorder_inorder_helper(vector<int> &preorder, int pre_start, int pre_end,
                                             vector<int> &inorder, int in_start, int in_end){
    if(in_start > in_end){
        return nullptr;
    }
    
    TreeNode *root;
    int i;
    //Find the root node in inorder list
    for(i = in_start; i <= in_end; i++)
        if(inorder[i] == preorder[pre_start])
            break;
    root = new TreeNode(inorder[i]);
    if(i > in_start)
        root->left = build_tree_preorder_inorder_helper(preorder, pre_start+1, pre_start+i-in_start, inorder, in_start, i-1);
    if(i < in_end)
        root->right = build_tree_preorder_inorder_helper(preorder, pre_start + i - in_start + 1, pre_end, inorder, i+1, in_end);
    
    return root;
}
/*
 * func: build_tree_preorder_inorder
 * goal: build the tree from the preorder and inorder traversal of the tree
 * @param preorder: preorder list of the tree
 * @param inorder: inorder list of the tree
 * return: root node of the tree
 */
TreeNode *build_tree_preorder_inorder(vector<int> &preorder, vector<int> &inorder){
    if(preorder.size() == 0){
        return nullptr;
    }
    
    return build_tree_preorder_inorder_helper(preorder, 0, preorder.size()-1, inorder, 0, inorder.size()-1);
}

Python Code:

# func: construct binary tree from preoder and inorder traversal
# @param preorder: preorder traversal list
# @param inorder: inorder traversal list
def build_tree_from_preorder_inorder(preorder, inorder):
    if not preorder:
        return None
    root = TreeNode(preorder[0])
    id = inorder.index(preorder[0])
    root.left = build_tree_from_preorder_inorder(preorder[1:id+1], inorder[0:id])
    root.right = build_tree_from_preorder_inorder(preorder[id+1:], inorder[id+1:])

    return root