[LeetCode 315] Count of Smaller Numbers After Self

作者 Shilei Tian 日期 2017-09-02
[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].

解题思路

这道题目实际上就是要求出每个元素的逆序数是多少。在 GeeksforGeeks 上面有一篇文章讲述如何利用归并排序来求一个数组逆序数的和,可以参见这里。那么我们如何改进让其能用到这里,为每个元素都求呢?我们看一个例子。

设我们有数组 arr = [6, 4, 1, 8, 7, 5, 2, 9],根据归并排序,我们需要将其划分为两个子数组:

6 4 1 8
7 5 2 9

这里我们为了说明原理,就不再重述对其进行拆分的过程。我们直接跳到归并这一步来,亦即我们现在两个子数组变成:

1 4 6 8
2 5 7 9

在往一起归并的时候,我们需要逐个判断来自第一个子数组和来自第二个子数组的元素。什么情况下会有逆序产生呢?在将第一个子数组的头元素插入到排序后数组的时候。什么意思呢?例如,我们将第一个子数组的头元素插入到已排序数组的时候,所有已经插入到已排序数组中的元素中,来自第二个子数组的元素均小于这个第一个子数组的头元素,那么当然我们就需要更新这个头元素对应的逆序数,令其加上所有已经插入到已排序数组中的元素中来自第二个子数组的元素的个数。这个值如何获得呢?其实很容易,因为我们在归并的时候,势必要维护两个标记元素 ij,分别指向第一个子数组和第二个子数组的头元素。i 应该初始化为本次归并的开始位置,j 应该初始化为这一段归并排序的中间位置 mid。已经插入到已排序数组中来自第二个数组的元素的个数,实际上就等于 j-mid。这个问题一下子就解决了。

那么下一个问题是,我们归并排序要修改元素的位置,那么这个怎么让我们能够关联到它最开始的下标呢?因为我们需要用这个最初的下标来定位记录逆序数的数组。这里有一个小技巧,我们无须对原数组进行排序,我们可以对它们的下标进行排序,亦即每次变换位置的,都是下标。我们维护这样一个下标的数组,就既可以访问到对应的元素,还可以访问到对应的逆序数数组。

最终的 C++ 代码如下:

class Solution {
private:
void mergeSort(const vector<int>& nums, vector<int>& indices, const int l, const int r, vector<int>& count) {
if (r - l > 1) {
int mid = l + (r - l) / 2;
mergeSort(nums, indices, l, mid, count);
mergeSort(nums, indices, mid, r, count);
vector<int> tmp;
tmp.reserve(r - l);
int i = l, j = mid;
while (i < mid && j < r) {
if (nums[indices[i]] <= nums[indices[j]]) {
tmp.push_back(indices[i]);
count[indices[i]] += j - mid;
++i;
} else {
tmp.push_back(indices[j]);
++j;
}
}
while (i < mid) {
tmp.push_back(indices[i]);
count[indices[i]] += j - mid;
++i;
}
while (j < r) {
tmp.push_back(indices[j]);
++j;
}
move(tmp.begin(), tmp.end(), indices.begin() + l);
}
}
public:
vector<int> countSmaller(vector<int>& nums) {
vector<int> ret(nums.size(), 0), indices(nums.size(), 0);
iota(indices.begin(), indices.end(), 0);
mergeSort(nums, indices, 0, nums.size(), ret);
return ret;
}
};

时间复杂度 O(nlogn),空间复杂度 O(nlogn),截止到目前为止超过 52.91% 的人。

其实我们可以对这段代码进行代码级的优化,就是让那个 tmp 数组不是每次递归进行分配和释放的,而是提前分配好,这样的话,时间效率会提升不少。但算法的层次已经够了。