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

首先想到的naive暴力，肯定会超时。代码如下：

class Solution {

    public List<Integer> countSmaller(int[] nums) {
        List<Integer> list = new ArrayList<>(nums.length);
        for (int i = 0; i < nums.length; i++) {
            list.add(i, getCount(i, nums));
        }
        return list;
    }

    private int getCount(int i, int[] nums) {
        int j = i + 1, count = 0;
        while (j < nums.length) {
            if (nums[j++] < nums[i]) count++;
        }
        return count;
    }

}

后来想到可用分治来做，分治可以考虑两种，一种是像QuickSort那样的分治，比如找第k大的数。而最多的是，归并排序这种框架的分治。下面是分治的代码：

class Solution {
    private List<Integer> list = null;

    public List<Integer> countSmaller(int[] nums) {
        list = new ArrayList<>(nums.length);
        for (int i = 0; i < nums.length; i++) list.add(0);
        Node[] nodes = getNodeList(nums);
        mergeSortWhileCount(nodes, 0, nodes.length - 1);
        return list;
    }

    private void mergeSortWhileCount(Node[] nodes, int s, int e) {
        if (s >= e) return;
        int mid = (s + e) / 2;
        mergeSortWhileCount(nodes, s, mid);
        mergeSortWhileCount(nodes, mid + 1, e);
        for (int i = s, j = mid + 1; i <= mid; i++) {
            while (j <= e && nodes[j].val < nodes[i].val) j++;
            int originIndex = nodes[i].index;
            list.set(originIndex, list.get(originIndex) + (j - mid - 1));
        }
        Arrays.sort(nodes, s, e + 1, Comparator.comparingInt(x -> x.val));
    }

    private Node[] getNodeList(int[] nums) {
        Node[] res = new Node[nums.length];
        for (int i = 0; i < nums.length; i++)
            res[i] = new Node(nums[i], i);
        return res;
    }


    static class Node {
        int val;
        int index;

        public Node(int val, int index) {
            this.val = val;
            this.index = index;
        }
    }
}

唯一思考的地方就是，由于归并排序会把之前的数组顺序打乱，导致没办法填充结果到相应的位置。所以需要记录一下之前的下标。

下面是通过二叉搜索树来做的(BST),开阔一下思路了。

class Solution {
    public List<Integer> countSmaller(int[] nums) {
        Integer[] ans = new Integer[nums.length];
        Node root = null;
        for (int i = nums.length - 1; i >= 0; i--) {
            root = insert(root, new Node(nums[i]), ans, i, 0);
        }
        return Arrays.asList(ans);
    }

    private Node insert(Node root, Node node, Integer[] ans, int i, int preSum) {
        if (root == null) {
            ans[i] = preSum;
            return node;
        } else if (root.val == node.val) {
            ans[i] = preSum + root.sum;
            root.dup++;
        } else if (node.val < root.val) {
            root.sum++;
            root.left = insert(root.left, node, ans, i, preSum);
        } else {
            root.right = insert(root.right, node, ans, i, preSum + root.sum + root.dup);
        }
        return root;
    }

    static class Node {
        int sum = 0, val, dup = 1;
        Node left, right;

        public Node(int val) {
            this.val = val;
        }
    }

}