Coder's Cat

LeetCode: Unique Binary Search Trees II

2020-04-01

How to generate all the unique binary search trees with given values [1, N].

Challenge Description

We iterate each element in the array, use current element(suppose as k) as root node, construct left subtrees with [1, k-1], and construct right subtrees with [k+1, N].

This is a depth first search algorithm:

class Solution {
public:
vector<TreeNode*> generateTrees(int n) {
vector<TreeNode*> empty;
return n == 0 ? empty : dfs(1, n);
}

vector<TreeNode*> dfs(int begin, int end){
vector<TreeNode*> res;
if(begin > end){
res.push_back(NULL);
return res;
}

if(begin == end){
TreeNode* pre = new TreeNode(begin);
res.push_back(pre);
return res;
}

for(int i=begin; i <= end; i++){
vector<TreeNode*> left = dfs(begin, i-1);
vector<TreeNode*> right = dfs(i+1, end);
for(int m = 0; m < left.size(); m++){
for(int n = 0; n < right.size(); n++){
TreeNode* pre = new TreeNode(i);
pre->left = left[m];
pre->right = right[n];
res.push_back(pre);
}
}
}
return res;
}
};

Optimization with memo

When we iterate the ranges, there are some duplicated computations could be cached.

For example, dfs(1, 2) is computed when 3 is root, it will be computed again when 4 is root.

So we could add a map to cache the result:

public:
map<pair<int,int>, vector<TreeNode*> > hmap;
vector<TreeNode*> generateTrees(int n) {
vector<TreeNode*> empty;
return n == 0 ? empty : dfs(1, n);
}

vector<TreeNode*> dfs(int begin, int end){
vector<TreeNode*> res;
pair<int, int> key(begin, end);
if(hmap.count(key))
return hmap[key];

if(begin > end){
res.push_back(NULL);
return res;
}

if(begin == end){
TreeNode* pre = new TreeNode(begin);
res.push_back(pre);
hmap[pair<int, int>(begin, end)] = res;
return res;
}

for(int i=begin; i <= end; i++){
vector<TreeNode*> left = dfs(begin, i-1);
vector<TreeNode*> right = dfs(i+1, end);
for(int m = 0; m < left.size(); m++){
for(int n = 0; n < right.size(); n++){
TreeNode* pre = new TreeNode(i);
pre->left = left[m];
pre->right = right[n];
res.push_back(pre);
}
}
}
hmap[key] = res;
return res;
}
};

Preparing for an interview? Check out this!