LeetCode 315. Count of Smaller Numbers After Self

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

分析

从题目中可以知道，我们的任务是快速统计每个数字右边的比它小的数的数字数量，既然如此，二叉树当然是个不错的选择。所以我们可以：

• 维护一颗树
• 每个节点记录下其所有子节点的数量
• 插入节点的同时，返回比它小的节点的数量

代码实现

class Solution
{
public:
    class Tree
    {
    public:
        //小于该节点的节点
        Tree *left;
        //大于或等于该节点的节点
        Tree *right;
        int value;
        int num;
        Tree(int v)
        {
            left = nullptr;
            right = nullptr;
            value = v;
            num = 0;
        }
        //插入节点并返回比它小的元素数量
        int insert(int v)
        {
            num += 1;
            if (value <= v)
            {
                int mod = 1;
                if (value == v)mod = 0;
                if (right == nullptr && left == nullptr)
                {
                    right = new Tree(v);
                    return mod;
                }
                else if (right == nullptr)
                {
                    right = new Tree(v);
                    return mod + 1 + left->num;
                }
                else if (left == nullptr)
                {
                    return mod + right->insert(v);
                }
                else
                {
                    return mod + 1 + right->insert(v) + left->num;
                }
            }
            else
            {
                if ( left == nullptr)
                {
                    left = new Tree(v);
                    return 0;
                }
                else
                {
                    return left->insert(v);
                }
            }
        }
    };
    vector<int> countSmaller(vector<int>& nums) {
        vector<int> result(nums.size(), 0);
        if (!nums.size())return result;
        Tree root(nums.back());
        for (int index = nums.size() - 2; index >= 0; index--)
        {
            result[index] = root.insert(nums[index]);
        }
        return result;
    }
};