Kth Largest Element in an Array

NeetCode: Kth Largest Element in an Array · LeetCode: 215. Kth Largest Element in an Array · Difficulty: Medium

Return the k-th largest element of an unsorted array (the k-th in sorted order, not the k-th distinct value).

Example

Text

Input:  nums = [2, 3, 1, 5, 4], k = 2      Output: 4
Input:  nums = [2, 3, 1, 2, 4, 2], k = 4   Output: 2

Brute force: sort

Arrays.sort on an int[] sorts ascending, so the k-th largest is at index n - k. The code sorts a clone so the caller's array is not reordered.

Java

import java.util.Arrays;

class Solution {
    public int findKthLargest(int[] nums, int k) {
        int[] sorted = nums.clone();
        Arrays.sort(sorted); // ascending
        return sorted[sorted.length - k];
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 5, 4}, 2));
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 2, 4, 2}, 4));
    }
}

Output

4
2

Time: O(n log n). Space: O(n).

Better: a min-heap of size k

Same idea as Kth Largest Element in a Stream: keep the k largest values in a min-heap and read its top. Java's PriorityQueue is a min-heap by default.

Java

import java.util.PriorityQueue;

class Solution {
    public int findKthLargest(int[] nums, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>(); // min-heap
        for (int x : nums) {
            heap.add(x);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        return heap.peek();
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 5, 4}, 2));
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 2, 4, 2}, 4));
    }
}

Output

4
2

Time: O(n log k). Space: O(k).

Optimal on average: quickselect

The k-th largest is the element at index n - k once sorted. Quickselect partitions like quicksort (the code works on a copy of the array), but recurses into only the side that contains the target index, so the expected work is n + n/2 + n/4 + ... = O(n). The pivot is the middle element, which handles already-sorted input well.

Java

class Solution {
    private int partition(int[] nums, int lo, int hi) {
        int mid = lo + (hi - lo) / 2;
        swap(nums, mid, hi);
        int pivot = nums[hi];
        int store = lo;
        for (int i = lo; i < hi; i++) {
            if (nums[i] < pivot) {
                swap(nums, i, store);
                store++;
            }
        }
        swap(nums, store, hi);
        return store;
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }

    public int findKthLargest(int[] nums, int k) {
        int[] a = nums.clone();
        int target = a.length - k;
        int lo = 0, hi = a.length - 1;
        while (true) {
            int p = partition(a, lo, hi);
            if (p == target) {
                return a[p];
            } else if (p < target) {
                lo = p + 1;
            } else {
                hi = p - 1;
            }
        }
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 5, 4}, 2));
        System.out.println(s.findKthLargest(new int[]{2, 3, 1, 2, 4, 2}, 4));
        System.out.println(s.findKthLargest(new int[]{7}, 1));
    }
}

Output

4
2
7

Time: O(n) on average, O(n²) in the worst case. Space: O(1) extra; the array is partitioned in place.

The worst case occurs when every pivot is close to the minimum or maximum. A random pivot makes that unlikely. The partition above also degrades to O(n²) when every element is equal, because no element is smaller than the pivot and each pass removes only one element; the heap version's O(n log k) is the guaranteed bound.