CountOfSmallerNumbersAfterSelf.java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class CountOfSmallerNumbersAfterSelf {
    static class TreeNode {
        int val;
        int count = 1;
        TreeNode left;
        TreeNode right;
        TreeNode(int x) { val = x; }
    }

    static int insertNode(TreeNode root, int val) {
        int thisCount = 0;
        while (true) {
            if (val <= root.val) {
                root.count++;
                if (root.left == null) {
                    root.left = new TreeNode(val);
                    break;
                } else {
                    root = root.left;
                }
            } else {
                thisCount += root.count;
                if (root.right == null) {
                    root.right = new TreeNode(val);
                    break;
                } else {
                    root = root.right;
                }
            }
        }
        return thisCount;
    }

    static List<Integer> countSmaller(int[] nums) {
        List<Integer> result = new ArrayList();
        if (nums == null || nums.length == 0) { return result; }
        int n = nums.length;
        TreeNode root = new TreeNode(nums[n - 1]);
        result.add(0);
        for (int i = n - 2; i >= 0; i--) {
            int count = insertNode(root, nums[i]);
            result.add(count);
        }
        Collections.reverse(result);
        return result;
    }

    public static void main(String[] args) {
        int[] nums = {5, 2, 6, 1};
        System.out.println(countSmaller(nums));
        int[] nums2 = {5, 2, 6, 4, 1, 3};
        System.out.println(countSmaller(nums2));
    }
}