Count of Smaller Numbers After S

题目
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].

答案

错误的答案(未考虑duplicate elements的情况)

class Solution {
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> counts = Arrays.asList(new Integer[nums.length]);
        TreeSet<Integer> tset = new TreeSet<>();
        int n = nums.length;
        for(int i = n - 1; i >= 0; i--) {
            TreeSet<Integer> subset = (TreeSet<Integer>)tset.headSet(nums[i]);
            counts.set(i, subset.size());
            tset.add(nums[i]);
        }
        return counts;
    }
}

正确

class Solution {
    class TreeNode {
        int val;
        int same_val_cnt = 1;
        int less_than_cnt = 0;
        TreeNode left, right;
        public TreeNode(int val) {
            this.val = val;
            left = right = null;
        }
    } 
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> counts = Arrays.asList(new Integer[nums.length]);
        TreeSet<Integer> tset = new TreeSet<>();
        int n = nums.length;
        int[] ans = new int[1];
        TreeNode root = null;
        for(int i = n - 1; i >= 0; i--) {
            ans[0] = 0;
            root = insert(root, nums[i], ans);
            counts.set(i, ans[0]);
        }
        return counts;
    }
    
    private TreeNode insert(TreeNode root, int val, int[] ans) {
        if(root == null) return new TreeNode(val);
        if(val == root.val) {
            root.same_val_cnt++;
            ans[0] += root.less_than_cnt;
        }
        else if(val < root.val) {
            root.less_than_cnt++;
            root.left = insert(root.left, val, ans);
        }
        else {
            ans[0] += root.same_val_cnt + root.less_than_cnt;
            root.right = insert(root.right, val, ans);
        }
        return root;
    }
}