C++ 9 lines O(n log n) map, plus stack with binary search

  • 4

    We populate right subtree if numbers are decreasing. If the current number is larger than any numbers in the right subtree, we find the first that is smaller, and make it a left subtree of the current number. The current number becomes a leaf of the right subtree.

    I use map as it provides the convenient insert operation, that returns the position of the inserted element in O (log n).

    TreeNode* constructMaximumBinaryTree(vector<int>& nums) {
        map<int, TreeNode*> q = { { nums[0], new TreeNode(nums[0]) } };
        for (auto i = 1; i < nums.size(); ++i) {
            auto it_b = q.insert({ nums[i], new TreeNode(nums[i]) });
            if (it_b.first != q.begin()) { // left tree.
                it_b.first->second->left = next(it_b.first, -1)->second;
                q.erase(q.begin(), it_b.first);
            if (next(it_b.first, 1) != q.end()) // right tree.
                next(it_b.first, 1)->second->right = it_b.first->second;
        return q.rbegin()->second;

    As pointed out by @Mrsuyi, we can use stack instead of map, as map adds some overhead to maintain internal BST. Below is the solution based on the one proposed by @Mrsuyi, optimized to use the binary search (as our stack will be always sorted).

    TreeNode* constructMaximumBinaryTree(vector<int>& nums) {
        vector<TreeNode*> s { new TreeNode(nums[0]) };
        for (int i = 1; i < nums.size(); ++i) {
            TreeNode* cur = new TreeNode(nums[i]);
            auto it = upper_bound(s.rbegin(), s.rend(), cur, 
                                  [](const TreeNode* a, const TreeNode* b) { return a->val < b->val; });
            if (it != s.rend()) (*it)->right = cur;
            if (it != s.rbegin()) cur->left = *next(it, -1);
            s.resize(distance(it, s.rend()));
        return s.front();

  • 1

    Brilliant idea !
    I optimized this solution slightly by using a stack and I hope you don't mind my posting it here
    Thanks !

  • 0

    @Mrsuyi good idea to use stack, I was considering it too but it was a little trickier to implement. In your solution, you could use binary search (as your stack is naturally sorted) instead of popping elements one by one. When you find a larger element, you can use vector.resize() to quickly remove all smaller elements. rbegin() and rend() can simplify the logic a bit, as the stack will be sorted in the descending order.

    I added another solution to my answer based on your code, optimized using the binary search.

    It looks like LeetCode OJ does not have a good test case to see the difference, so I ran both algorithms using a large random input, and the version with the binary search was 20% faster.

    My initial algorithm that uses set is several times slower. I am guessing that this is due to the overhead to maintain BST internally.

  • 0

    Can somebody explain this solution in simple words?

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.