LeetCode 315. 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].

Main Idea:
  Convert the array into BST, smaller number is the number on the left tree.  Time complexity O(NlgN)

#include <vector>
#include <iostream>
using namespace std;

struct TreeNode {
  TreeNode* left;
  TreeNode* right;
  int val;
  TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

void insertToBST(TreeNode* root, int num) {
  TreeNode* tmp = root;
  TreeNode* parent;
  while(tmp) {
    parent = tmp;
    if(tmp->val < num) {
      tmp = tmp->right;
    } else {
      tmp = tmp->left;
    }
  }
  TreeNode* tmp_1 = new TreeNode(num);
  if(num > parent->val) parent->right = tmp_1;
  else parent->left = tmp_1;
}

TreeNode* leftMost(TreeNode* root) {
  TreeNode* parent = root;
  TreeNode* tmp = root;
  while(tmp) {
    parent = tmp;
    tmp = tmp->left;
  }
  return parent;
}
TreeNode* convertToBST(vector<int>& nums) {
  TreeNode* root = new TreeNode(nums[0]);
  for(int i = 1; i < nums.size(); ++i) {
    insertToBST(root, nums[i]);
  }
  return root;
}

int countNumber(TreeNode* root) {
  if(!root) return 0;
  return 1 + countNumber(root->left) + countNumber(root->right);
}

void printTree(TreeNode* root) {
  if(root == NULL) return;
  printTree(root->left);
  cout << root->val << endl;
  printTree(root->right);
}

int main(void) {
  vector<int> nums{5,2,6,1};
  TreeNode* root = convertToBST(nums);
//  printTree(root);
  int i = 0;
  // time complexity: O(Nlg(N))
  while(i < nums.size()) {
    vector<int> tmp(nums.begin() + i, nums.end());
    TreeNode* root = convertToBST(tmp);
    cout << countNumber(root->left) << endl;
    i++;
  }
}