[LeetCode] Kth Largest Element in an Array

标签：

Well, this problem has a naive solution, which is to sort the array in descending order and return the k-1-th element. However, sorting algorithm gives O(nlogn) complexity, which may cause TLE. Indeed, this problem has an O(n) solution.

Well, the O(n) solution has no mystery. It is also closely related to sorting. Specifically, it uses the partition method of quicksort.

In quicksort, in each iteration, we need to select a pivot and then partition the array into three parts:

1. elements smaller than the pivot;
2. elements equal to the pivot;
3. elements larger than the pivot.

Now, let‘s do an example taken the array [3, 2, 1, 5, 4, 6] in the problem statement. Let‘s assume in each time we select the leftmost element to be the pivot, in this case, 3. We then use it to partition the array into the above 3 parts, which results in [2, 1, 3, 5, 4, 6]. Now 3 is in the third position and we know that it is the third smallest element. Now, do you realize that this step can be used to solve this problem?

In fact, the above partition puts elements smaller than the pivot before the pivot and thus the pivot will then be the kth smallest element if it is at k-1. Since the problem requires us to find the kth largest element, we can simply modify the partition to put elements larger than the pivot before the pivot. That is, after partition, the array becomes [5, 4, 6, 3, 2, 1]. Now we know that 3 is the4-th largest element. If we are asked to find the 2-th largest element, then we know it is left to3. If we are asked to find the 5-th largest element, then we know it is right to 3. So, the problem is reduced to approximately half of its original size, giving the recursion T(n) = T(n/2) + O(n) in which O(n) is the time for partition. This recursion, once solved, gives T(n) = O(n) and thus we have a linear time solution. Note that since we only need to consider one half of the array, the time complexity is O(n). If we need to consider both the two halves of the array, like in quicksort, then the recursion will be T(n) = 2T(n/2) + O(n) and the complexity will be O(nlogn).

Let‘s formally write down the algorithm above.

1. Initialize left to be 0 and right to be nums.size() - 1;
2. Partition the array, if the pivot is at the k-1-th position, returns it (we are done);
3. If the pivot is right to the k-1-th position, update right to be the left neighbor of the pivot;
4. Else update left to be the right neighbor of the povot.
5. Repeat 2.

Now that we have the above algorithm, let‘s turn it into code.

 1 int partition(vector<int>& nums, int left, int right) {
 2     int pivot = nums[left];
 3     int l = left + 1, r = right;
 4     while (l <= r) {
 5         if (nums[l] < pivot && nums[r] > pivot)
 6             swap(nums[l++], nums[r--]);
 7         else if (nums[l] >= pivot) l++;
 8         else if (nums[r] <= pivot) r--;
 9     }
10     swap(nums[left], nums[r]);
11     return r;
12 }
13 int findKthLargest(vector<int>& nums, int k) {
14     int left = 0, right = nums.size() - 1;
15     while (true) {
16         int pos = partition(nums, left, right);
17         if (pos == k - 1) return nums[pos];
18         if (pos < k - 1) left = pos + 1;
19         else right = pos - 1;
20     }
21 }

[LeetCode] Kth Largest Element in an Array

标签：

原文地址：http://www.cnblogs.com/jcliBlogger/p/4548045.html