ByteInTheSky

Kth Largest Element in an Array

Tagged: Array, Divide and Conquer, Sorting, Priority Queue, Heap, Quickselect

Problem

Given an array of integers, our task is to find Kth Largest Element in the Array. The Kth largest element here is in the sorted order, not in the Kth distinct element.
For example,

Input: nums = [3,2,1,5,6,4], k = 2
Output: 5
Explanation: 5 is the second largest element in the sorted order.

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4
Explanation: 4 is the fourth largest element in the sorted order.

Solution

Sorting

fun findKthLargest(nums: IntArray, k: Int): Int {
    nums.sort()
    return nums[nums.size - k]
}

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

Priority Queue (Heap) - Two Pass

fun findKthLargest(nums: IntArray, k: Int): Int {
    val priorityQueue = PriorityQueue(
        compareByDescending<Int> { it }
    )
    nums.forEach {
        priorityQueue.add(it)
    }

    var count = 1
    while (priorityQueue.isNotEmpty()) {
        val num = priorityQueue.poll()
        if (count == k)
            return num

        count++
    }

    return 0
}

public int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(
            Comparator.reverseOrder()
    );
    for (int num : nums) {
        priorityQueue.add(num);
    }

    var count = 1;
    while (!priorityQueue.isEmpty()) {
        int num = priorityQueue.poll();
        if (count == k)
            return num;

        count++;
    }

    return 0;
}

Priority Queue (Heap) - One Pass

fun findKthLargest(nums: IntArray, k: Int): Int {
    val priorityQueue = PriorityQueue(
        compareBy<Int> { it }
    )
    nums.forEach {
        if (priorityQueue.size < k)
            priorityQueue.add(it)
        else {
            if (priorityQueue.peek() < it) {
                priorityQueue.remove()
                priorityQueue.add(it)
            }
        }
    }

    return priorityQueue.peek()
}

public int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();
    for (int num : nums) {
        if (priorityQueue.size() < k)
            priorityQueue.add(num);
        else {
            if (priorityQueue.peek() < num) {
                priorityQueue.remove();
                priorityQueue.add(num);
            }
        }
    }

    return priorityQueue.peek();
}

Quickselect

fun findKthLargest(nums: IntArray, k: Int): Int {
    return quickSelect(nums, k, 0, nums.size - 1) 
}

private fun quickSelect(nums: IntArray, k: Int, low: Int, high: Int): Int {
    var l = low
    var h = high
    val smallestElementIndex = nums.size - k
    while (l <= h) {
        if (l == h)
            return nums[l]
        val random = Random()
        var pivotIndex = l + random.nextInt(h - l)
        pivotIndex = partition(nums, l, h, pivotIndex)
        if (pivotIndex == smallestElementIndex)
            return nums[pivotIndex]
        else if (smallestElementIndex < pivotIndex)
            h = pivotIndex - 1
        else
            l = pivotIndex + 1
    }

    return -1
}

private fun partition(nums: IntArray, low: Int, high: Int, pivotIndex: Int): Int {
    val pivot = nums[pivotIndex]

    nums[pivotIndex] = nums[high].also {
        nums[high] = nums[pivotIndex]
    }

    var storeIndex = low

    for (i in low until high) {
        if (nums[i] < pivot) {
            nums[storeIndex] = nums[i].also {
                nums[i] = nums[storeIndex]
            }

            storeIndex++
        }
    }

    nums[storeIndex] = nums[high].also {
        nums[high] = nums[storeIndex]
    }

    return storeIndex
}

public int findKthLargest(int[] nums, int k) {
    return quickSelect(nums, k, 0, nums.length - 1);
}

private int quickSelect(int[] nums, int k, int low, int high) {
    int smallestElementIndex = nums.length - k;
    while (low <= high) {
        if (low == high)
            return nums[low];
        Random random = new Random();
        var pivotIndex = low + random.nextInt(high - low);
        pivotIndex = partition(nums, low, high, pivotIndex);
        if (pivotIndex == smallestElementIndex)
            return nums[pivotIndex];
        else if (smallestElementIndex < pivotIndex)
            high = pivotIndex - 1;
        else
            low = pivotIndex + 1;
    }

    return -1;
}

private int partition(int[] nums, int low, int high, int pivotIndex) {
    int pivot = nums[pivotIndex];

    swap(nums, high, pivotIndex);

    int storeIndex = low;

    for (int i = low; i < high; i++) {
        if (nums[i] < pivot) {
            swap(nums, storeIndex, i);
            storeIndex++;
        }
    }

    swap(nums, storeIndex, high);

    return storeIndex;
}

private void swap(int[] nums, int a, int b) {
    int temp = nums[b];
    nums[b] = nums[a];
    nums[a] = temp;
}
← Insert Delete GetRandom O(1)
Top K Frequent Elements →