See all coding puzzles and patterns here.


Algorithm Description

Quick select is an adaptation of the quicksort algorithm to find the top or bottom $k$ elements in a collection. The algorithm entails doing a quicksort until we randomly find that a pivot that sits at index $k$, which indicates that all items before that item are less than it and all items after that item are greater than it.


Performance Characteristics

Let $n = $ the number of items in the input collection.

Since quicksort has a worse case runtime of $O(n^2)$, the worst case of quick select is also $O(n^2)$, since we might randomly pick every other pivot before we pick $k$.

However, the average case of quick select should be better, between $log(n)$ and $n$.


Example Implementation

import random

def quick_select(nums: list[int], k: int) -> list[int]:
    """
    Returns the lowest k numbers in nums.
    """
    n = len(nums)
    if n < k:
        return nums
    def partition(start: int, end: int) -> int:
        pivot_idx = random.randint(start, end)
        pivot = nums[pivot_idx]
        nums[end], nums[pivot_idx] = nums[pivot_idx], nums[end]
        l, r = 0, end
        while l < r:
            while l < r and nums[l] < pivot:
                l += 1
            while l < r and nums[r] >= pivot:
                r -= 1
            nums[l], nums[r] = nums[r], nums[l]
        nums[l], nums[end] = nums[end], nums[l]
        return l
    def helper(start: int, end: int):
        if start > end:
            return
        p = partition(start, end)
        if p == k:
            return
        if k < p:
            helper(start, p)
        else:
            helper(p + 1, end)
    helper(0, n - 1)
    return nums[:k]

Puzzles

Puzzle Difficulty External Link
K Closest Points to Origin Medium Leetcode 973