LintCode-Kth Largest Element

Find K-th largest element in an array.

Note

You can swap elements in the array

Example

In array [9,3,2,4,8], the 3th largest element is 4

Challenge

O(n) time, O(1) space

Analysis:

We use Quick partition to perform the selection. We use the method of "median of medians", i.e., median of every group of 5 elements, to select the pivot. In this way, we can guarantee that the complexity in the worst case is still O(n).

Reference:

http://blog.csdn.net/v_JULY_v/article/details/6370650

http://blog.csdn.net/v_july_v/article/details/6431001

Solution:

class Solution {
    //param k : description of k
    //param numbers : array of numbers
    //return: description of return
    public int kthLargestElement(int k, ArrayList<Integer> numbers) {
        return selectK(numbers,0,numbers.size()-1,k);
    }

    public void insertSort(ArrayList<Integer> numbers, int start, int end){
        for (int i=start;i<=end;i++){
            int key = numbers.get(i);
            int ind = i-1;
            while (ind>=start && numbers.get(ind)<key){
                numbers.set(ind+1,numbers.get(ind));
                ind--;
            }
            numbers.set(ind+1,key);
        }
    }        

    public void swap(ArrayList<Integer> numbers, int x, int y){
        int temp = numbers.get(x);
        numbers.set(x,numbers.get(y));
        numbers.set(y,temp);
    }

    //This function selects the kth largest number in the range numbers[start...end].
    //For example, if start=8, end=11, k=2, we will select the 2nd largest number in the range numbers[8...11].
    //So in the code, we need be very carefull about the index of the kth largest number.
    public int selectK(ArrayList<Integer> numbers, int start, int end, int k){
        if (end-start+1<=5){
            insertSort(numbers,start,end);
            return numbers.get(start+k-1);
        }

        //Get the median of every group of 5 elements.
        int pos = start;
        int cur = start;
        while (true){
            int curEnd = Math.min(cur+4,end);
            //For even number, we pick up the later median, that is why there is plus one at last.
            int median5 = selectK(numbers,cur,curEnd,(curEnd-cur+1)/2+1);
            swap(numbers,pos,(curEnd+cur+1)/2);
            pos++;
            if (curEnd==end) break;
            else cur = curEnd+1;
        }

        //Get the median of median5s, i.e., median of [start...pos-1].
        int median = selectK(numbers,start,pos-1,(pos-start)/2+1);
        int mInd = start;
        for (int i=start;i<=start+4;i++) 
            if (numbers.get(i)==median){
                mInd = i;
                break;
            }

        
        //Use the median as the pivot to partition [start...end].
        swap(numbers,mInd,end);
        int p1 = start, p2 = end-1;
        while (p1<end && numbers.get(p1)>median) p1++;
        while (p2>0 && numbers.get(p2)<median) p2--;
        while (p1<p2){
            swap(numbers,p1,p2);
            p1++;
            p2--;
            while (p1<end && numbers.get(p1)>median) p1++;
            while (p2>0 && numbers.get(p2)<median) p2--;
        }
        swap(numbers,p1,end);
        if (p1-start+1==k)
            return numbers.get(p1);
        else if (p1-start+1>k)
            return selectK(numbers,start,p1-1,k);
        else 
            return selectK(numbers,p1+1,end,k-(p1-start+1));
    }
            
        
            
};