Sorting & Quickselect

Custom comparators, quickselect, counting and bucket sort, and merge-sort counting.

You'll almost never be asked to implement a library sort from scratch. But sorting ideas are everywhere in interviews: custom comparators, quickselect for the kth element, counting and bucket sort for O(n) solutions, merge sort for counting inversions, and the Dutch national flag partition. This lesson covers the ones that come up.

Know the built-in sort

Python's sort and sorted use Timsort: O(n log n), stable, and O(n) on input that's already nearly sorted.
Stable means equal elements keep their original order. It lets you sort by several keys in separate passes, and it matters whenever ties have to stay in input order.
Sort by several keys with a tuple, negating numbers for descending order: key=lambda x: (-x.score, x.name).
For comparisons that can't be expressed as a key, like *Largest Number*, use functools.cmp_to_key.

from functools import cmp_to_key

def largest_number(nums):
    strs = sorted(map(str, nums), key=cmp_to_key(lambda a, b: (b + a > a + b) - (b + a < a + b)))
    return "0" if strs[0] == "0" else "".join(strs)

Quickselect: the kth element in O(n) on average

Partition around a pivot, as in quicksort, but recurse only into the side that contains k. With a random pivot, the average time is O(n). The worst case is O(n²).

import random

def find_kth_largest(nums, k):
    target = len(nums) - k                       # index in ascending order
    lo, hi = 0, len(nums) - 1
    while True:
        pivot = nums[random.randint(lo, hi)]
        lt, i, gt = lo, lo, hi                   # three-way partition handles duplicates
        while i <= gt:
            if nums[i] < pivot:
                nums[lt], nums[i] = nums[i], nums[lt]
                lt += 1
                i += 1
            elif nums[i] > pivot:
                nums[i], nums[gt] = nums[gt], nums[i]
                gt -= 1
            else:
                i += 1
        if target < lt:
            hi = lt - 1
        elif target > gt:
            lo = gt + 1
        else:
            return pivot

[!INTERVIEW] Heap or quickselect? A min-heap of size k is O(n log k), simple, and works on a stream. Quickselect is O(n) on average, but it modifies the array and has a bad worst case. Offer the heap first, then mention quickselect as the faster option. That's usually the ideal answer.

Dutch national flag: three-way partition

Sort an array of 0s, 1s and 2s in one pass, using the same partition as quickselect above. Everything before lt is 0, everything after gt is 2, and i scans the middle. It's O(n) time, O(1) space, and a frequent warm-up question.

Counting and bucket sort: beating O(n log n)

When the values come from a small range or can be bucketed, sorting doesn't need comparisons:

Counting sort: count each value, then rebuild. O(n + range).
Bucket by frequency: for *Top K Frequent Elements*, put each value in a bucket indexed by its count, from 0 to n, then read the buckets from the highest down. O(n).
Pigeonhole: *Maximum Gap* and *First Missing Positive* (place each value at index value − 1) use the same idea.

from collections import Counter

def top_k_frequent(nums, k):
    buckets = [[] for _ in range(len(nums) + 1)]
    for val, freq in Counter(nums).items():
        buckets[freq].append(val)
    out = []
    for freq in range(len(buckets) - 1, 0, -1):
        out += buckets[freq]
        if len(out) >= k:
            return out[:k]
    return out

Merge sort: when you need the merge step

Merge sort is the algorithm to implement when asked to "sort an array" with guaranteed O(n log n), or to sort a linked list. Its merge step also counts pairs across the two halves. That solves inversion counting, *Reverse Pairs* and *Count of Smaller Numbers After Self* in O(n log n).

def sort_and_count(a):
    if len(a) <= 1:
        return a, 0
    mid = len(a) // 2
    left, x = sort_and_count(a[:mid])
    right, y = sort_and_count(a[mid:])
    merged, i, j, inv = [], 0, 0, x + y
    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            merged.append(left[i]); i += 1
        else:
            merged.append(right[j]); j += 1
            inv += len(left) - i          # every remaining left element is > right[j]
    merged += left[i:] + right[j:]
    return merged, inv

When sorting is the whole trick

Many problems become easy once the input is sorted: intervals, two pointers for 3Sum, greedy scheduling, and grouping anagrams by their sorted letters. When you're stuck, ask yourself: "What if the input were sorted?" Sorting costs O(n log n), which is usually acceptable.

[!TIP] Practice Kth Largest Element · Sort Colors · Top K Frequent Elements · Sort an Array · Largest Number · K Closest Points · Sort Characters by Frequency · Reverse Pairs