[LeetCode] 315、计算右侧小于当前元素的个数

题目描述

给定一个整数数组 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 个更小的元素.

解题思路

第一反应是暴力法，这肯定可以做，但是暴力法肯定对不起这个Hard题。这个题和“计算逆序对（归并排序）算法”有相似之处。这部分一定要看大佬文章。（这个题有点难理解，可以暂时略过）

给出几个参考文章：

• 计算右侧小于当前元素的个数“的六种方法
• 树状数组（Python 代码、Java 代码）
• 二分查找（$C++$）

参考代码

基于“归并排序”，略复杂，不是很推荐

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

public class Solution {

    private int[] temp;
    private int[] counter;
    private int[] indexes;

    public List<Integer> countSmaller(int[] nums) {
        List<Integer> res = new ArrayList<>();
        int len = nums.length;
        if (len == 0) {
            return res;
        }
        temp = new int[len];
        counter = new int[len];
        indexes = new int[len];
        for (int i = 0; i < len; i++) {
            indexes[i] = i;
        }
        mergeAndCountSmaller(nums, 0, len - 1);
        for (int i = 0; i < len; i++) {
            res.add(counter[i]);
        }
        return res;
    }

    /**
     * 针对数组 nums 指定的区间 [l, r] 进行归并排序，在排序的过程中完成统计任务
     *
     * @param nums
     * @param l
     * @param r
     */
    private void mergeAndCountSmaller(int[] nums, int l, int r) {
        if (l == r) {
            // 数组只有一个元素的时候，没有比较，不统计
            return;
        }
        int mid = l + (r - l) / 2;
        mergeAndCountSmaller(nums, l, mid);
        mergeAndCountSmaller(nums, mid + 1, r);
        // 归并排序的优化，同样适用于该问题
        // 如果索引数组有序，就没有必要再继续计算了
        if (nums[indexes[mid]] > nums[indexes[mid + 1]]) {
            mergeOfTwoSortedArrAndCountSmaller(nums, l, mid, r);
        }
    }

    /**
     * [l, mid] 是排好序的
     * [mid + 1, r] 是排好序的
     *
     * @param nums
     * @param l
     * @param mid
     * @param r
     */
    private void mergeOfTwoSortedArrAndCountSmaller(int[] nums, int l, int mid, int r) {
        // 3,4  1,2
        for (int i = l; i <= r; i++) {
            temp[i] = indexes[i];
        }
        int i = l;
        int j = mid + 1;
        // 左边出列的时候，计数
        for (int k = l; k <= r; k++) {
            if (i > mid) {
                indexes[k] = temp[j];
                j++;
            } else if (j > r) {
                indexes[k] = temp[i];
                i++;
                // 此时 j 用完了，[7,8,9 | 1,2,3]
                // 之前的数就和后面的区间长度构成逆序
                counter[indexes[k]] += (r - mid);
            } else if (nums[temp[i]] <= nums[temp[j]]) {  // 这里必须为<=
                indexes[k] = temp[i];
                i++;
                // 此时 [4,5, 6   | 1,2,3 10 12 13]
                //           mid          j
                counter[indexes[k]] += (j - mid - 1);
            } else {
                // nums[indexes[i]] > nums[indexes[j]] 构成逆序
                indexes[k] = temp[j];
                j++;
            }
        }
    }

    public static void main(String[] args) {
        int[] nums = new int[]{5, 2, 6, 1};
        Solution solution = new Solution();
        List<Integer> countSmaller = solution.countSmaller(nums);
        System.out.println(countSmaller);
    }
}

利用二叉排序树（跪着看完，推荐！）

struct BSTNode{
    int val;
    int cnt;  // cnt代表的次数也就是在nums数组中比val小的数的个数
    BSTNode *left;
    BSTNode *right;
    BSTNode(int x): val(x), left(nullptr), right(nullptr), cnt(0){
        // nothing
    }
};

void insert2Tree(BSTNode *root, BSTNode *insert_node, int &count_small){
    // 这里并不是严格的“二叉搜索树”，因为可能存在相等元素
    // 注意小于等于插入到左子树，防止多加1（记住，比如 6, 5, 5）
    if(insert_node->val <= root->val){
        // 插入的结点更小,被比较结点(即node)的count++，然后插入到左子树(如果不为空)
        root->cnt++;

        if(root->left)
            insert2Tree(root->left, insert_node, count_small);
        else  // 左子树为空，插入结点就作为当前结点的左孩子
            root->left = insert_node;
    }else{
        // 插入的结点更大，需要在右子树(如果不为空)继续找
        count_small += root->cnt + 1;

        if(root->right)
            insert2Tree(root->right, insert_node, count_small);
        else  // 当前右子树为空，插入结点作为当前结点右孩子
            root->right = insert_node;
    }
}
/*count_small作为一个引用的参数，在递归寻找子树的时候作为一个“类似全局变量”的存在*/


class Solution {
public:
    vector<int> countSmaller(vector<int>& nums) {
        int n=nums.size();
        // 如果数组为空返回空值
        if(n==0)
            return {};

        vector<int> res;
        res.push_back(0);

        // 建立一个二叉搜素树
        BSTNode* root = new BSTNode(nums[n-1]);
        int count_small;
        for(int i = n - 2; i >= 0; i--){
            count_small = 0;
            insert2Tree(root, new BSTNode(nums[i]), count_small);
            res.push_back(count_small);
        }

        // 最后不要忘记删除树节点
        delete root;
        // push_back的时候是逆序的，此时只要将count数组reverse即可
        reverse(res.begin(), res.end());
        return res;
    }
};

更易理解的版本（递归的精髓是：管好当下，之后的事抛给递归）

class Solution {
public:
    // binary search tree
    struct Node {
        int val, smaller;  // 这里的 smaller 其实是某元素右侧“小于等于”当前元素的个数
        Node *left, *right;
        Node(int v, int s) : val(v), smaller(s), left(nullptr), right(nullptr) {}
    };
    
    int insert(Node *&root, int val) {  // 注意这里是 Node *&root
        if (!root) {
            root = new Node(val, 0);
            return 0;
        }
        if (val <= root->val) {
            root->smaller++;
            return insert(root->left, val);
        } else {
            return insert(root->right, val) + root->smaller + 1;
        }
    }
    
    vector<int> countSmaller(vector<int>& nums) {
        vector<int> res(nums.size());
        Node *root = nullptr;
        for (int i = nums.size() - 1; i >= 0; --i) {
            res[i] = insert(root, nums[i]);
        }
        return res;
    }
};

字节面试题：计算右侧大于当前元素的个数

#include<bits/stdc++.h>
using namespace std;

struct TreeNode{
    int val;
    int max_num;
    struct TreeNode *left, *right;
    TreeNode(int a, int b): val(a), max_num(b), left(nullptr), right(nullptr){}
};

int insert(TreeNode *& root, int val){
    if(!root){
        root = new TreeNode(val, 0);
        return 0;
    }

    if(val < root->val){
        return insert(root->left, val) + root->max_num + 1;
    }else{
        root->max_num++;
        return insert(root->right, val);
    }
}


int main(){
    vector<int> vec = {2, 5, 1, 3, 4};
    int len = vec.size();
    TreeNode *root = nullptr;

    vector<int> ans(len);
    for(int i = len-1; i >= 0; i--){
        ans[i] = insert(root, vec[i]);
    }

    for(auto num: ans){
        cout << num << ' ';
    }

    return 0;
}