[LeetCode] 315. Count of Smaller Numbers After Self

Problem

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:

Input: [5,2,6,1]
Output: [2,1,1,0]
Explanation:
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.

Solution

class Solution {
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> res = new ArrayList<>();
        if (nums == null || nums.length == 0) return res;
        int len = nums.length;
        TreeNode root = new TreeNode(nums[len-1]);
        res.add(0);
        for (int i = len-2; i >= 0; i--) {
            int count = insertNode(root, nums[i]);
            res.add(count);
        }
        Collections.reverse(res);
        return res;
    }
    
    private int insertNode(TreeNode root, int num) {
        int count = 0; //the # of smaller numbers of num
        while (root != null) {
            if (num <= root.val) {
                root.count++;
                if (root.left == null) {
                    root.left = new TreeNode(num);
                    break;
                } else root = root.left;
            } else {
                count += root.count;
                if (root.right == null) {
                    root.right = new TreeNode(num);
                    break;
                } else root = root.right;
            }
        }
        return count;
    }
}
class TreeNode {
    int val;
    int count = 1;
    TreeNode left;
    TreeNode right;
    public TreeNode(int val) {
        this.val = val;
    }
}