原题

LintCode 5. Kth Largest Element

Description

Find K-th largest element in an array.

Notice

You can swap elements in the array

Example

In array [9,3,2,4,8], the 3rd largest element is 4.

In array [1,2,3,4,5], the 1st largest element is 5, 2nd largest element is 4, 3rd largest element is 3 and etc.

解题

题意是从数组中找到第n大的数。

这里我们实现寻找第n个最小的数，那么第n个最大的数其实就是第nums.size() - n + 1个最小的数。

最简单的思路就是，将整个数组排序，O(nlogn)的时间复杂度下问题可以得到解决。但这个过程一定是有优化空间的，因为排序的时候我们做了很多无意义的操作：假设最后第n小的数是x，我们只关心有多少个值比x小，丝毫不关心这些值的顺序，更不关心比x大的数之间的顺序。

下面尝试以分治的思想解决这个问题：
首先我们随机地在数组中取一个值pivot，然后将数组分为三个部分：

• 比pivot小的部分 SL，长度为L
• 与pivot相等的部分 SM，长度为M
• 比pivot大的部分 SR，长度为R

我们要找目标——第n小的数字，有下列三种情况：

• n <= L ： 则目标是SL中的第n小的数
• L < n <= L + M ： 则目标就是pivot
• L + M < n ： 则目标是SR中的第(n - L - M)小的值

此时只要将上述过程递归进行下去，就可以得到结果。

代码

class Solution {
public:
    /*
    * @param n: An integer
    * @param nums: An array
    * @return: the Kth largest element
    */
    int kthLargestElement(int n, vector<int> &nums) {
        // write your code here
        return kthSmallestElement(nums.size() - n + 1, nums, 0, nums.size() - 1);
    }
private:
    int kthSmallestElement(int n, vector<int> &nums, int begin, int end) {
        // write your code here
        if (begin == end) return nums[begin];
        int pivot = nums[rand() % (end - begin + 1) + begin];
        int left = begin;
        for (int i = left; i <= end; i++) {
            if (nums[i] < pivot) swap(nums[i], nums[left++]);
        }
        int right = left;
        for (int i = right; i <= end; i++) {
            if (nums[i] == pivot) swap(nums[i], nums[right++]);
        }

        int leftLength = left - begin;
        int middleLength = right - left;

        if (n <= leftLength) {
            return kthSmallestElement(n, nums, begin, begin + leftLength - 1);
        } else if (leftLength < n && n <= leftLength + middleLength) {
            return pivot;
        } else {
            return kthSmallestElement(n - leftLength - middleLength, nums, right, end);
        }
    }
};