Jack Li's Blog

0215.Kth Largest Element in an Array

class Solution {
public:
    int partition(vector<int>& nums, int left, int right) {
        // choose pivot
        int pivot = nums[right];
        int l = left;
        int r = right - 1;

        // swap left and right based on pivot
        while(l <= r) {
            if(nums[l] < pivot && nums[r] > pivot) swap(nums[l++], nums[r--]);
            if(nums[l] >= pivot) l++;
            if(nums[r] <= pivot) r--;
        }
        // 5  3 4
        // lr
        // 5 3 4
        // r l
        // swap pivot to l
        swap(nums[right], nums[l]);

        return l;
    }

    int findKthLargest(vector<int>& nums, int k) {
        // use quickselect
        // Time: O(n)
        // Space: O(1)

        int l = 0;
        int r = nums.size() - 1;
        k = k - 1; // convert to index-0

        while(l <= r) {
            int idx = partition(nums, l, r);
            if(idx == k) {
                return nums[idx];
            } else if (idx > k) {
                r = idx - 1;
            } else {
                l = idx + 1;
            }
        }

        return -1;
    }
};