Kth Largest Element in an Array (Quickselect)

Problem

Find the kth largest element in an unsorted array (not necessarily the kth distinct element), better than full sorting on average, using Quickselect.

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

Python Solution

import random

def find_kth_largest(nums, k):
    target_index = len(nums) - k

    def partition(left, right, pivot_index):
        pivot = nums[pivot_index]
        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]
        store_index = left
        for i in range(left, right):
            if nums[i] < pivot:
                nums[store_index], nums[i] = nums[i], nums[store_index]
                store_index += 1
        nums[right], nums[store_index] = nums[store_index], nums[right]
        return store_index

    left, right = 0, len(nums) - 1
    while left < right:
        pivot_index = random.randint(left, right)
        pivot_index = partition(left, right, pivot_index)
        if pivot_index == target_index:
            break
        elif pivot_index < target_index:
            left = pivot_index + 1
        else:
            right = pivot_index - 1
    return nums[target_index]

C++ Solution

int partition(vector<int>& nums, int left, int right, int pivotIndex) {
    int pivot = nums[pivotIndex];
    swap(nums[pivotIndex], nums[right]);
    int storeIndex = left;
    for (int i = left; i < right; i++) {
        if (nums[i] < pivot) swap(nums[storeIndex++], nums[i]);
    }
    swap(nums[right], nums[storeIndex]);
    return storeIndex;
}

int findKthLargest(vector<int>& nums, int k) {
    int targetIndex = nums.size() - k;
    int left = 0, right = nums.size() - 1;
    while (left < right) {
        int pivotIndex = left + rand() % (right - left + 1);
        pivotIndex = partition(nums, left, right, pivotIndex);
        if (pivotIndex == targetIndex) break;
        else if (pivotIndex < targetIndex) left = pivotIndex + 1;
        else right = pivotIndex - 1;
    }
    return nums[targetIndex];
}

Java Solution

class Solution {
    public int findKthLargest(int[] nums, int k) {
        int targetIndex = nums.length - k;
        int left = 0, right = nums.length - 1;
        Random random = new Random();
        while (left < right) {
            int pivotIndex = left + random.nextInt(right - left + 1);
            pivotIndex = partition(nums, left, right, pivotIndex);
            if (pivotIndex == targetIndex) break;
            else if (pivotIndex < targetIndex) left = pivotIndex + 1;
            else right = pivotIndex - 1;
        }
        return nums[targetIndex];
    }

    private int partition(int[] nums, int left, int right, int pivotIndex) {
        int pivot = nums[pivotIndex];
        int temp = nums[pivotIndex]; nums[pivotIndex] = nums[right]; nums[right] = temp;
        int storeIndex = left;
        for (int i = left; i < right; i++) {
            if (nums[i] < pivot) {
                int t = nums[storeIndex]; nums[storeIndex] = nums[i]; nums[i] = t;
                storeIndex++;
            }
        }
        int t = nums[right]; nums[right] = nums[storeIndex]; nums[storeIndex] = t;
        return storeIndex;
    }
}

Complexity

Time: O(n) on average, since each partition step discards roughly half the remaining range, though O(n²) in the rare worst case of consistently poor pivot choices; random pivot selection makes this practically very unlikely. Space: O(1) — partitioning happens in place, with no extra array required.

Written by Lorens Mishra

Software Engineer Notes Management System Administrator

Discussion

0 comments
Sign in to join the discussion.

No comments yet. Be the first to say something.