Update 0215-kth-largest-element-in-an-array.cpp

Author: unw9527

Diff for: cpp/0215-kth-largest-element-in-an-array.cpp

@@ -69,27 +69,44 @@ class Solution {
 */
 
 // Video's QuickSelect implementation
+// class Solution {
+// public:
+//     int findKthLargest(vector<int>& nums, int k) {
+//         int index = nums.size() - k;
+//         return quickSelect(nums, index, 0, nums.size() - 1);
+//     }
+// private:
+//     int quickSelect(vector<int>& nums, int k, int l, int r){
+//         int pivot = nums[r];
+//         int p_pos = l;
+//         for (int i = l; i < r; ++i){
+//             if (nums[i] <= pivot){
+//                 swap(nums[i], nums[p_pos]);
+//                 ++p_pos;
+//             }
+//         }
+//         swap(nums[p_pos], nums[r]);
+//         if (k < p_pos)
+//             return quickSelect(nums, k, l, p_pos - 1);
+//         if (k > p_pos)
+//             return quickSelect(nums, k, p_pos + 1, r);
+//         return nums[p_pos];
+//     }
+// };
+
+// Solution that can pass current tests
 class Solution {
 public:
     int findKthLargest(vector<int>& nums, int k) {
-        int index = nums.size() - k;
-        return quickSelect(nums, index, 0, nums.size() - 1);
-    }
-private:
-    int quickSelect(vector<int>& nums, int k, int l, int r){
-        int pivot = nums[r];
-        int p_pos = l;
-        for (int i = l; i < r; ++i){
-            if (nums[i] <= pivot){
-                swap(nums[i], nums[p_pos]);
-                ++p_pos;
-            }
+        int pivot = nums[rand() % nums.size()];
+        vector<int> left, mid, right;
+        for (auto num: nums) {
+            if (num > pivot) left.push_back(num);
+            else if (num == pivot) mid.push_back(num);
+            else right.push_back(num);
         }
-        swap(nums[p_pos], nums[r]);
-        if (k < p_pos)
-            return quickSelect(nums, k, l, p_pos - 1);
-        if (k > p_pos)
-            return quickSelect(nums, k, p_pos + 1, r);
-        return nums[p_pos];
+        if (k <= left.size()) return findKthLargest(left, k);
+        else if (k <= left.size() + mid.size()) return mid[0];
+        else return findKthLargest(right, k - left.size() - mid.size());
     }
 };