Java刷题笔记11： 计算右侧小于当前元素的个数

题目描述

给定一个整数数组 nums，按要求返回一个新数组 counts。数组 counts 有该性质： counts[i] 的值是 nums[i]
右侧小于 nums[i] 的元素的数量。

示例

输入: [5,2,6,1]
输出: [2,1,1,0]
解释:
5 的右侧有 2 个更小的元素 (2 和 1).
2 的右侧仅有 1个更小的元素 (1).
6 的右侧有 1 个更小的元素 (1).
1 的右侧有 0 个更小的元素.

题目链接

插入排序+二分查找

逆序遍历nums，用一个数组list从小到大存储目前遍历过的元素，那么在每遍历到一个nums里的元素nums[i]时，都可用二分查找的方式在list里找到比nums[i]小的元素数量,然后将这个元素插入list，由于是逆序遍历的原因，所以list里的元素一定是num[i]右边的元素。

public static List<Integer> countSmaller(int[] nums) {
        List<Integer> list = new ArrayList<>(); //保存nums内从小到大的元素
        List<Integer> ans = new LinkedList<>();
        for (int i = nums.length - 1; i >= 0; i--) {
            int minNum = binarySearchMin(list, nums[i]);
            ans.add(0,minNum);
        }
        return ans;
    }
    private static int binarySearchMin(List<Integer> list, int num) {
        int left = 0;
        int right = list.size()-1;
        while (left <= right) {
            int mid =( (right - left) >> 1) + left;
            if (list.get(mid) < num && (mid+1>=list.size()||list.get(mid + 1) >= num )) {
                list.add(mid+1,num);
                return mid+1;
            } else if (list.get(mid) >= num) {
                right = mid - 1;
            } else left = mid + 1;
        }
        list.add(0,num);
        return 0;
    }

以上的代码实际上时间复杂度还是较高的。主要是由于二分查找后的插入操作时间复杂度较高。因为list底层是数组，所以在插入时需要移动后面的元素，导致插入操作的最坏时间复杂度达到了O(N）。为了优化时间复杂度，我们可以将list换成二叉搜索树。这时插入便不在需要额外的时间了，在搜索时便已经找到了num[i]需要插入的位置，所以搜索和插入这两个操作的时间复杂度为Log(N).总的时间复杂度为O(nlogn）.

二叉搜索树

public List<Integer> countSmaller(int[] nums) {
        //初始化
        Integer[] res = new Integer[nums.length];
        Arrays.fill(res, 0);
        List<Integer> list = new ArrayList<>();
        //反向构造二叉树，统计右边的较小数
        TreeNode root = null;
        for (int i = nums.length - 1; i >= 0; i--){
            root = addAndCount(root, new TreeNode(nums[i]), res, i);
        }
        return Arrays.asList(res);
}

public TreeNode addAndCount(TreeNode root, TreeNode node, Integer[] res, int i){
        if(root == null){
            root = node;
            return root;
        }
        //根节点的左边保存不大于根节点的元素
        if(root.val >= node.val){ 
            //统计左节点的元素个数
            root.count++;
            root.left = addAndCount(root.left, node, res, i);
        }else{
            //走到右边获取该元素左边的个数（根节点 1 + 左节点 root.count）
            res[i] += 1 + root.count;
            //统计右边是否还有更小的元素
            root.right = addAndCount(root.right, node, res, i);
        }
        return root;
    }
}

class TreeNode{
    int val;
    int count;
    TreeNode left, right;

    public TreeNode(int val){
        this.val = val;
        this.count = 0;
        left = null;
        right = null;
    }
}

这段代码引用自LeetCode用户UFO的题解。因为懒再加上二叉搜索树还是比较好实现的。关键在于需要在每个节点保存比它小的数的数量，也就是左子树的节点数。若不保存那么计算数量时时间复杂度就会变高，本人一开始也就是因为没想到这点，感觉数组对于计算较小数的数量更方便所以采取了数组的方式。很感谢这篇题解的作者指点迷津。