Thursday, May 22, 2014

LeetCode: Unique Binary Search Trees II

Given n, generate all structurally unique BST's (binary search trees) that store values 1...n.

For example,
Given n = 3, your program should return all 5 unique BST's shown below.

   1         3     3      2      1
    \       /     /      / \      \
     3     2     1      1   3      2
    /     /       \                 \
   2     1         2                 3


Following the previous problem, for every possible root node, we first generate all possible left subtrees and all possible right subtrees, then create the root node and push every combination to the final set.

C++ Code:

/*
 * func: generate_unique_BST_helper
 * goal: helper function to generate BSTs
 * @param start: start number of the tree
 * @param end: end number of the tree
 * return: all possible trees
 */
vector<TreeNode *> generate_unique_BST_helper(int start, int end){
    vector<TreeNode *> result;
    if(start > end){
        result.emplace_back(nullptr);
        return result;
    }
    for(int i = start; i <= end; ++i){
        vector<TreeNode *> left_subtree = generate_unique_BST_helper(start, i-1);
        
        vector<TreeNode *> right_subtree = generate_unique_BST_helper(i + 1, end);
        
        for(TreeNode* left : left_subtree){
            for(TreeNode* right : right_subtree){
                TreeNode *root = new TreeNode(i);
                root->left = left;
                root->right = right;
                result.emplace_back(root);
            }
        }
    }
    return result;
}
/*
 * func: generate_unique_BST
 * goal: generate all unique binary search trees by given number
 * @param n: number of nodes
 * return: all possible BST
 */
vector<TreeNode *> generate_unique_BST(int n){
    return generate_unique_BST_helper(1, n);
}

Python Code:

# func: generate all possible BSTs
# @param n: total number of nodes
# @return: all possible trees
def generate_unique_BSTs(n):
    def generate_BST(start, end):
        if start > end:
            return [None]

        result = []
        for i in xrange(start, end+1):
            #Generate all possible left subtree
            left_subtree = generate_BST(start, i-1)
            #Generate all possible right subtree
            right_subtree = generate_BST(i+1, end)

            #Create root node and push every combinations to the result
            for left in left_subtree:
                for right in right_subtree:
                    root = TreeNode(i)
                    root.left = left
                    root.right = right
                    result.append(root)
        return result

    return generate_BST(1, n)

No comments:

Post a Comment