Leetcode: 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].

Solution with BST:  

Time: O(NlogN), space: O(N) better than naive solution O(N^2)

TreeNode的结构稍微改变了一下, 加入一个count值，表示该Node左子树有多少个node（包括自己）

public class Solution {
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> res = new ArrayList<Integer>();
        if (nums == null || nums.length == 0) return res;
        TreeNode root = new TreeNode(nums[nums.length-1]);
        res.add(0);
        for (int i=nums.length-2; i>=0; i--) {
            res.add(insert(root, nums[i]));
        }
        Collections.reverse(res);
        return res;
    }
    
    public int insert(TreeNode root, int num) {
        int thisCount = 0;
        while (true) {
            if (root.val >= num) {
                root.count++;
                if (root.left == null) {
                    root.left = new TreeNode(num);
                    break;
                }
                else {
                    root = root.left;
                }
            }
            else {
                thisCount += root.count;
                if (root.right == null) {
                    root.right = new TreeNode(num);
                    break;
                }
                else {
                    root = root.right;
                }
            }
        }
        return thisCount;
    }
    
    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;
        int count;
        public TreeNode(int value) {
            this.val = value;
            this.count = 1;
        }
    }
}

 我更习惯另一种写法：TreeNode含有count(自己有几个重复的), leftSize（自己左子树有几个node）

第8行可以用res.add(0, insert()), 这样不用reverse，但是好像比较慢

public class Solution {
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> res = new ArrayList<Integer>();
        if (nums == null || nums.length == 0) return res;
        TreeNode root = new TreeNode(nums[nums.length-1]);
        res.add(0);
        for (int i=nums.length-2; i>=0; i--) {
            res.add(insert(root, nums[i]));
        }
        Collections.reverse(res);
        return res;
    }
    
    public int insert(TreeNode root, int num) {
        int thisCount = 0;
        while (true) {
            if (root.val > num) {
                root.leftSize++;
                if (root.left == null) {
                    root.left = new TreeNode(num);
                    break;
                }
                else {
                    root = root.left;
                }
            }
            else if (root.val == num) {
                thisCount += root.leftSize;
                root.count++;
                break;
            }
            else {
                thisCount = thisCount + root.leftSize + root.count;
                if (root.right == null) {
                    root.right = new TreeNode(num);
                    break;
                }
                else {
                    root = root.right;
                }
            }
        }
        return thisCount;
    }
    
    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;
        int count;
        int leftSize;
        public TreeNode(int value) {
            this.val = value;
            this.count = 1;
            this.leftSize = 0;
        }
    }
}