LeetCode: Unique Binary Search Trees II

#leetcode #algorithms

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

Challenge Description

Depth-first search

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;
    }
};