215. Kth Largest Element in an Array

Published on May 6, 2023

Question

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

You must solve it in O(n) time complexity.

Algorithm

Use priorityQueue: By construct a priority queue, we poll elements until next one is the nth largest one.

Code1

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        for (int num : nums) {
            pq.offer(num);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        return pq.poll();
    }
}

Algorithm2

QuickSelect Algorithm

This is an awesome algorithm. The Wiki gives us a very detailed explanation, you could that one. Here, I only explain the code based on this particular question 215.

Two main for this answer if you want to use the quick select algorithm.

Quick select
- By using a pivot(similar thoughts in quick sort), each time you do a select, the pivot's position is the final position, and the elements before it are all smaller that the nums[pivot], and the number behind it it are all larger(assume we are constructing a nums from small to large). So if you get a pivot smaller than k, you need to find the larger part of nums. Otherwise, go to the smaller side.
Partition
- The is the procedure that moves numbers based on the pivot. In the end, we got the position where pivot should be at, and we swap the number of that position and the pivot we set at first.
- The algorithm chooses the left most number as pivot each time. Some choose right most, some choose random, and some choose the median-of-3. The choose of pivot may affect the time complexity of the algorithm.
In the quick select, we do it in place, the space complexity is O(n) for we are using the stack in recursion. And the time complicity is O(n) (but the worst time is O(n^2) . Why the average time complexity is O(n)

Code2-0

class Solution {
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, nums.length - k);
    }
    
    private int quickSelect(int[] nums, int left, int right, int k) {
        if (left == right) return nums[left];
        int pivot = partition(nums, left, right);
        if (pivot == k) {
            return nums[k];
        } else if (pivot <= k) {
            return quickSelect(nums, pivot + 1, right, k);           
        } else {
            return quickSelect(nums, left, pivot - 1, k);  
        }
    }
    
    private int partition(int[] nums, int left, int right) {
        int pivot = nums[left];
        int l = left + 1;
        int r = right;
        while (l <= r) {
            if (nums[r] < pivot && nums[l] > pivot) {
                swap(nums, l, r);
            } 
            if (nums[l] <= pivot) {
                l++;
            }
            if (nums[r] >= pivot) {
                r--;
            }
        }
        swap(nums, left, r);
        return r;
    }
    
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

Well you could modify slightly to remove the tail recursion. Using a loop.

Code2-1

class Solution {
    public int findKthLargest(int[] nums, int k) {
      
        int l = 0, r = nums.length - 1;
        while (true) {
            int pos = partition(nums, l, r);
            if(nums.length - pos == k) {
                return nums[pos];
            } else if (nums.length - pos < k) {
                r = pos - 1;
            } else {
                l = pos + 1;
            }
        }
    }
    
    private int partition(int[] nums, int left, int right) {
        int pivot = nums[left];
        int l = left + 1;
        int r = right;
        while (l <= r) {
            if (nums[r] < pivot && nums[l] > pivot) {
                swap(nums, l, r);
            } 
            if (nums[l] <= pivot) {
                l++;
            }
            if (nums[r] >= pivot) {
                r--;
            }
        }
        swap(nums, left, r);
        return r;
    }
    
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

Or you could sort the algorithm from large to small so you don't have to check the position nums.length - k.

Code2-2

class Solution {
    public int findKthLargest(int[] nums, int k) {
        if(nums == null || nums.length == 0) return 0;
        int l = 0; 
        int r = nums.length-1; 
        
        while(true) {
            int pos = partition(nums, l, r);
            if(pos + 1 == k) {
                return nums[pos];
            } else if (pos + 1 > k) {
                r = pos - 1;
            } else {
                l = pos + 1;
            }
        }
    }
    private int partition(int[] nums, int left, int right) {
        int pivotIndex = left;
        int pivot = nums[pivotIndex];
        int l = left + 1;
        int r = right;
        while(l <= r) {
            if(nums[l] < pivot && nums[r] > pivot) {
                swap(nums, l++, r--);
            }
            if(nums[l] >= pivot) l++;
            if(nums[r] <= pivot) r--;
        }
        swap(nums, pivotIndex, r);
        return r;
    }
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}