[LeetCode315] Count of Smaller Numbers After Self

You are given an integer array nums and you have to return a new counts array. The counts array has the property where counts[i] is the number of smaller elements to the right of nums[i].

Example:

Given nums = [5, 2, 6, 1]

To the right of 5 there are 2 smaller elements (2 and 1).
To the right of 2 there is only 1 smaller element (1).
To the right of 6 there is 1 smaller element (1).
To the right of 1 there is 0 smaller element.
Return the array [2, 1, 1, 0].

Hide Company Tags Google

就用BST吧！O(nlogn) w/ O(n) spcae!

class Solution {
public:
    class TreeNode{
      public:
      int val, cnt;
      TreeNode *left, *right;
      TreeNode(int val, int cnt){
          this->val = val;
          this->cnt = cnt;
          this->left = nullptr;
          this->right = nullptr;
      }
    };
    vector<int> countSmaller(vector<int>& nums) {
        TreeNode* root = nullptr;
        vector<int> res(nums.size());
        for(int i = nums.size()-1; i>=0; --i){
            TreeNode* node = new TreeNode(nums[i], 0);
            root = insertNode(root, node);
            res[i] = query(root, nums[i]);
        }
        return res;
    }

    TreeNode* insertNode(TreeNode* root, TreeNode* node){
        if(!root) return node;
        TreeNode* cur = root;
        while(cur){
            if(cur->val > node->val){//insert to left subtree
                ++cur->cnt;
                if(!cur->left){
                    cur->left = node;
                    break;
                }else cur = cur->left;
            }else{//insert to right subtree;
                if(!cur->right){
                    cur->right = node;
                    break;
                }else cur = cur->right;
            }
        }
        return root;
    }
    int query(TreeNode* root, int target){
        if(!root) return 0;
        TreeNode* cur = root;
        int count = 0;
        while(cur){
            if(cur->val > target){
                cur = cur->left;
            }else if(cur->val < target) {
                count += cur->cnt + 1;
                cur = cur->right;
            }
            else return (cur->cnt) + count;
        }
        return 0;
    }
};