Select K'th smallest element

TopK.java

    /**
     * select k'th smallest element
     * @param nums the array to select
     * @param k 0 <= k < |nums|
     * @return k'th element
     */
    private int select(int[] nums, int k) {
        int lo = 0, hi = nums.length - 1;

        // If just one element, goto return directly;
        while (lo < hi) {
            int i = partition(nums, lo, hi);
            if (i > k) {
                hi = i - 1;
            } else if (i < k) {
                lo = i + 1;
            } else {
                return nums[i];
            }
        }

        return nums[lo];
    }
    
    /**
     * partition the subarray a[lo..hi] so that a[lo..j-1] <= a[j] <= a[j+1..hi] and return the index j.
     * @param nums the array to partition, 2 elements in the array at least
     * @param lo lo >= 0
     * @param hi hi <= |nums|
     * @return j
     */
    private int partition(int[] nums, int lo, int hi) {
        int x = nums[lo];

        int i = lo, j = hi + 1;

        while (true) {
            while (nums[++i] < x)
                if (i == hi)
                    break;
            while (nums[--j] > x)
                if (j == lo)
                    break;
            if (i >= j)
                break;
            int tmp = nums[i];
            nums[i] = nums[j];
            nums[j] = tmp;
        }

        int tmp = nums[j];
        nums[j] = x;
        nums[lo] = tmp;
        
        return j;
    }