Leetcode 215. Kth Largest Element in an Array

/**
 * How quick sort works.
 * First choose a pivot, say the first element of the array,
 * set two pointer low and high point to the start and end of the array,
 * from high to low, find a element that is samller than the pivot, replace nums[low] with nums[high]
 * then from low to high, find a element that is greater than the pivot, replace nums[high] with nums[low],
 * untill low == high,
 * set nums[low] as pivot, finish a partition.
 * Do the same for array [0, ..., index_of_pivot-1] and [index_of_pivot+1, ..., nums.length-1].
 * Therefore, index_of_pivot is the index_of_pivot largetest element of the array.
 * For this problem.
 * So, if k > index_of_pivot, we need to move to the left,
 * else if k < index_of_pivot, we need to move to the right,
 * else index_of_pivot is k. 
 */ 
public class Solution {
    public int findKthLargest(int[] nums, int k) {
        // note that pos is the (pos+1) samllest element in the array
        int low = 0, high = nums.length-1, pos = nums.length - k;
        
        // k is always valid
        while (true) {
            int idx = partition(nums, low, high);
            if (idx > pos) high = idx - 1;
            else if (idx < pos) low = idx + 1;
            else return nums[idx];
        }
    }
    
    // returns (low+1)th samllest element in the array
    public static int partition(int[] nums, int low, int high) {
        int pivot = nums[low];
        while (low < high) {
            // from high to low, find an elem that is less than the pivot
            while (low < high && nums[high] >= pivot) --high;
            nums[low] = nums[high];
            // from low to high, find an elem that is greater than the pivot
            while (low < high && nums[low] <= pivot) ++low;
            nums[high] = nums[low];
        }
        nums[low] = pivot;
        return low;
    }
}

AVG: T(n) = T(n/2) + O(n) = O(n)
WORST: T(n) = T(n-1) + O(n) = O(n^2)