Kth Largest Element in an Array

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

Input:
[3,2,1,5,6,4] and k = 2

Output:
 5

Example 2:

Input:
[3,2,3,1,2,4,5,5,6] and k = 4

Output:
 4

Note: You may assume k is always valid, 1 ≤ k ≤ array's length.

Thoughts:

1. naive: sort and get: O(N lg N) running time + O(1) memory

2. use min heap: O(N lg K) running time + O(K) memory

3. selection algorithm (partition method): O(NlogN) best case / O(N^2) worst case running time + O(1) memory

4. Code 1. Naive:

public int findKthLargest(int[] nums, int k) {
        final int N = nums.length;
        Arrays.sort(nums);
        return nums[N - k];
}

Code 2. Min Heap:

public int findKthLargest(int[] nums, int k) {

    final PriorityQueue<Integer> pq = new PriorityQueue<>();
    for(int val : nums) {
        pq.offer(val);

        if(pq.size() > k) {
            pq.poll();
        }
    }
    return pq.peek();
}

Code 2.1 Min Heap C++ (implementation: multiset)

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        multiset<int> min_heap; //min heap implemetation
        for (int num : nums){
            min_heap.insert(num);
            if(min_heap.size() > k) min_heap.erase(min_heap.begin());

        }

        return *min_heap.begin();
    }
};

Code 2.2 Min Heap: Python

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        pq = []
        for val in nums:
            heapq.heappush(pq, val)
            if len(pq) > k:
                heapq.heappop(pq)

        return heapq.heappop(pq)

Code 2.3 Max Heap: C++

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        priority_queue<int> pq(nums.begin(), nums.end());
        for (int i = 0; i < k - 1; i++){
            pq.pop(); // pop k - 1 times
        }

        return pq.top();
    }
};

Code 2.3 Max Heap C++ (implementation: max_heapify)

class Solution {
public:
    int findKthLargest(vector<int>& nums, int k) {
        build_max_heap(nums);
        // swap max element to the end k - 1 times, then kth max is nums[0]
        for (int i = 0; i < k - 1; i++){
            swap(nums[0], nums[heap_size - 1]);
            heap_size--;
            max_heapify(nums, 0);
        }

        return nums[0];
    }

private:
    int heap_size;

    inline int left_child(int idx){
        return (idx << 1) + 1;
    }

    inline int right_child(int idx){
        return (idx << 1) + 2;
    }

    void max_heapify(vector<int>& nums, int idx){
        int largest = idx, l = left_child(idx), r = right_child(idx);
        if (l < heap_size && nums[l] > nums[largest]) largest = l; 
        if (r < heap_size && nums[r] > nums[largest]) largest = r; 
        if (idx != largest){
            swap(nums[idx], nums[largest]);
            max_heapify(nums, largest); //percolate it down
        }
    }

    void build_max_heap(vector<int>&nums){
        heap_size = nums.size();
        for(int i = (heap_size >> 1) - 1; i>= 0; i--){
            max_heapify(nums, i);
        }  
    }
};

Code 3. Partition, QuickSelect:

class Solution {
    public int findKthLargest(int[] nums, int k) {
        k = nums.length - k; // we are finding N - k th smallest
        int low = 0, high = nums.length - 1;

        while(low < high){
            int i = quickSelect(nums, low, high);
            if(i < k){
                low = i + 1;
            }
            else if (i > k){
                high = i - 1;
            }
            else{
                break;
            }
        }

        return nums[k];
    }

    private int quickSelect(int [] nums, int low, int high){
        int i = low;
        int j = high + 1;

        while(true){
            while(i < high && nums[++i] < nums[low]);
            while(j > low && nums[low] < nums[--j]);
            if(i >= j) break;
            // swap nums[i] and nums[j]
            int tmp = nums[i];
            nums[i] = nums[j];
            nums[j] = tmp;

        }

        // swap pivot to the position j and return
        int tmp = nums[low];
        nums[low] = nums[j];
        nums[j] = tmp;
        return j;
    }
}

class Solution {
    public int findKthLargest(int[] nums, int k) {
        k = nums.length - k; // we are finding N - k + 1 th smallest: [N - (k - 1)] - 1
        int low = 0, high = nums.length - 1;

        while(low < high){
            int i = quickSelect(nums, low, high);
            // System.out.println(low + " "+ high + " " + i + " " + k);
            if(i < k){
                low = i + 1;
            }
            else if (i > k){
                high = i - 1;
            }
            else{
                break;
            }
        }

        return nums[k];
    }

    private int quickSelect(int [] nums, int low, int high){
        int i = low + 1;
        int j = high;
        // take low as the pivot
        while(true){
            while(i <= j && nums[i] <= nums[low]) i++;
            while(i <= j && nums[low] <= nums[j]) j--;
            if(i > j) break;
            // swap nums[i] and nums[j]
            int tmp = nums[i];
            nums[i] = nums[j];
            nums[j] = tmp;

        }

        // swap pivot to the position j and return
        int tmp = nums[low];
        nums[low] = nums[j];
        nums[j] = tmp;
        return j;
    }
}

Python:

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        def quickselect(nums, low, high):
            i, j = low + 1, high 
            while True:
                while i <= j and nums[i] <= nums[low]: i+= 1
                while j >= i and nums[j] >= nums[low]: j-= 1
                if i > j: break
                # swap nums[i], nums[j]
                nums[i], nums[j] = nums[j], nums[i]

            # swap nums[low], nums[j]
            nums[low], nums[j] = nums[j], nums[low]

            return j

        n, k = len(nums), len(nums) - k
        low, high = 0, n - 1

        while low < high:
            i = quickselect(nums, low, high)

            if i < k:
                low = i + 1

            elif i > k:
                high = i - 1

            else: break

        return nums[k]

So how can we improve the above solution and make it O(N) guaranteed (practically speaking)? The answer is quite simple, we can randomize the input, so that even when the worst case input would be provided the algorithm wouldn't be affected. (Blum-Floyd-Pratt-Rivest-Tarjan algorithm) In practice , when the array is long enough. The possibility is almost zero.