The k smallest and the median

Often you need only part of the order: the three fastest laps, the ten cheapest flights, the median salary. sorted(values)[:k] gives the k smallest, but it sorts the whole list, O(n log n), to keep only a few items. There are cheaper ways.

For one item, min(values) and max(values) go through the list once: O(n). A tempting way to get k items is to take the smallest, remove it, and repeat k times. Each round goes through the whole list again, so that's O(n·k): for the 5,000 smallest of 200,000 values, a billion steps.

heapq.nsmallest and nlargest

The module heapq has two functions made for this. heapq.nsmallest(k, values) returns the k smallest values from the smallest up, and heapq.nlargest(k, values) the k largest from the largest down. Both take key=, like sorted, and items with equal keys come in their original order:

import heapq

laps = [71.2, 68.9, 70.4, 69.5, 72.8, 68.9, 70.1]
print(heapq.nsmallest(3, laps))
print(heapq.nlargest(2, laps))

runners = [("Ann", 70.4), ("Bo", 68.9), ("Cy", 69.5), ("Di", 68.9)]
print(heapq.nsmallest(2, runners, key=lambda runner: runner[1]))

[68.9, 68.9, 69.5]
[72.8, 71.2]
[('Bo', 68.9), ('Di', 68.9)]

They go through the values once and keep the k best seen so far in a heap, a structure that finds and replaces the worst of those k quickly (the module on trees and heaps builds one). Each value costs at most O(log k), so O(n log k) in all, with O(k) extra memory. For a small k that's close to O(n).

Quickselect

To find the value that would be at index k after sorting, the k-th smallest counting from 0, reuse the partition from quicksort. After a partition, the pivot is at its final index p. If p is k, that's the answer. If k is smaller than p, the answer is in the part before the pivot, otherwise in the part after it, and only that part needs more work: the other part is never touched again. That's quickselect.

With good pivots each round halves the part, so the work is about n + n/2 + n/4 + ..., less than 2n: O(n) on average. Bad pivots every time make it O(n²), like quicksort, and the quicksort step's two cures apply here too. Take the middle item as the pivot (or a random one, which no order of the items makes slow every time). And use the three-way partition: with the two-way one, a long run of equal values would come back round after round, one item fewer each time, just as in quicksort. With three parts, all the items equal to the pivot are done in one round.

Since only one part goes on, a loop is enough, with no recursion and no depth limit to worry about:

def partition3(values, low, high, pivot):
    lt = low
    i = low
    gt = high
    while i <= gt:
        if values[i] < pivot:
            values[lt], values[i] = values[i], values[lt]
            lt += 1
            i += 1
        elif values[i] > pivot:
            values[i], values[gt] = values[gt], values[i]
            gt -= 1
        else:
            i += 1
    return lt, gt


def kth_smallest(values, k):
    values = values[:]
    low = 0
    high = len(values) - 1
    while True:
        pivot = values[(low + high) // 2]
        lt, gt = partition3(values, low, high, pivot)
        if k < lt:
            high = lt - 1
        elif k > gt:
            low = gt + 1
        else:
            return pivot


times = [71, 64, 90, 58, 77, 69, 83]
print(kth_smallest(times, 0), kth_smallest(times, 3), kth_smallest(times, 6))
print(times)

58 71 90
[71, 64, 90, 58, 77, 69, 83]

partition3 is the quicksort step's three-way partition, working only on the part from low to high. After it, the items equal to the pivot fill indexes lt to gt, which are their places in the sorted list. If k is among those indexes, the pivot is the answer. If k is smaller, only the part before lt goes on, and if it's bigger, only the part after gt. The pivot comes from the part itself, so at least one item is done each round and the loop ends.

Look at values = values[:]. Partitioning rearranges the list, and a function that quietly reorders its caller's list is a nasty surprise. The copy costs O(n), which keeps the whole O(n).

The median

The median is the middle of the sorted values: half of them are below it and half above. It tells you more about a typical salary or response time than the average, which a few huge values can pull far up. With an odd number of values n, it's the value at index n // 2 of the sorted list. With an even number, it's the average of the two in the middle, at n // 2 - 1 and n // 2. So the median of [40, 15, 25] is 25, and the median of [9, 30, 13, 20] is (13 + 20) / 2 = 16.5. One or two quickselects find it in O(n) on average, without sorting. (Python's statistics.median sorts the values: O(n log n).)