{"id":"dsa-sorting-k-smallest","slug":"k-smallest","title":"The k smallest and the median","text":"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.\n\nFor 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.\n\n### heapq.nsmallest and nlargest\n\nThe 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:\n\n```python\nimport heapq\n\nlaps = [71.2, 68.9, 70.4, 69.5, 72.8, 68.9, 70.1]\nprint(heapq.nsmallest(3, laps))\nprint(heapq.nlargest(2, laps))\n\nrunners = [(\"Ann\", 70.4), (\"Bo\", 68.9), (\"Cy\", 69.5), (\"Di\", 68.9)]\nprint(heapq.nsmallest(2, runners, key=lambda runner: runner[1]))\n```\n\n```output\n[68.9, 68.9, 69.5]\n[72.8, 71.2]\n[('Bo', 68.9), ('Di', 68.9)]\n```\n\nThey 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).\n\n### Quickselect\n\nTo 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**.\n\nWith 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.\n\nSince only one part goes on, a loop is enough, with no recursion and no depth limit to worry about:\n\n```python\ndef partition3(values, low, high, pivot):\n    lt = low\n    i = low\n    gt = high\n    while i <= gt:\n        if values[i] < pivot:\n            values[lt], values[i] = values[i], values[lt]\n            lt += 1\n            i += 1\n        elif values[i] > pivot:\n            values[i], values[gt] = values[gt], values[i]\n            gt -= 1\n        else:\n            i += 1\n    return lt, gt\n\n\ndef kth_smallest(values, k):\n    values = values[:]\n    low = 0\n    high = len(values) - 1\n    while True:\n        pivot = values[(low + high) // 2]\n        lt, gt = partition3(values, low, high, pivot)\n        if k < lt:\n            high = lt - 1\n        elif k > gt:\n            low = gt + 1\n        else:\n            return pivot\n\n\ntimes = [71, 64, 90, 58, 77, 69, 83]\nprint(kth_smallest(times, 0), kth_smallest(times, 3), kth_smallest(times, 6))\nprint(times)\n```\n\n```output\n58 71 90\n[71, 64, 90, 58, 77, 69, 83]\n```\n\n`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.\n\nLook 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).\n\n### The median\n\nThe **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).)\n","task":"A shop wants its three most expensive products, the most expensive first. Fill in the two blanks: the `heapq` function, and the index that makes the key a product's price. Products with equal prices keep their order. The program should print:\n\n```\ndesk 120\nrug 120\nchair 85\n```\n","kind":"fill","seed":"import heapq\n\nproducts = [(\"lamp\", 35), (\"desk\", 120), (\"pen\", 2), (\"chair\", 85), (\"rug\", 120)]\nfor name, price in heapq.[[nlargest]](3, products, key=lambda product: product[[[1]]]):\n    print(name, price)\n","solution":"import heapq\n\nproducts = [(\"lamp\", 35), (\"desk\", 120), (\"pen\", 2), (\"chair\", 85), (\"rug\", 120)]\nfor name, price in heapq.nlargest(3, products, key=lambda product: product[1]):\n    print(name, price)\n","hints":["`nsmallest` gives the smallest items and `nlargest` the largest. The key takes one product, a tuple, and returns its price.","The function is `nlargest`, and the price is the second part of each tuple, at index 1: `product[1]`."],"tests":[{"name":"Output","stdin":"","expect":"desk 120\nrug 120\nchair 85","hidden":false}],"require":[],"forbid":[],"more":[{"task":"A delivery app needs the k shops closest to the customer, who is at the point (0, 0). Write `k_closest(shops, k)`, where `shops` is a list of `(x, y)` tuples. Return a list of the k closest shops, the closest first, with shops at the same distance in the order they have in the list. The distance from (0, 0) to (x, y) is the square root of x² + y², but comparing x² + y² gives the same order without a square root. `k` is from 0 to the number of shops (for 0, return an empty list). It must be fast for the 5,000 closest of 200,000 shops.\n\nFor example, `k_closest([(3, 4), (1, 1), (-2, 0), (0, 5)], 2)` returns:\n\n```\n[(1, 1), (-2, 0)]\n```\n","kind":"code","seed":"def k_closest(shops, k):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(k_closest([(3, 4), (1, 1), (-2, 0), (0, 5)], 2))\n","solution":"import heapq\n\n\ndef k_closest(shops, k):\n    return heapq.nsmallest(k, shops, key=lambda shop: shop[0] * shop[0] + shop[1] * shop[1])\n\n\nif __name__ == \"__main__\":\n    print(k_closest([(3, 4), (1, 1), (-2, 0), (0, 5)], 2))\n","hints":["You need only k of the shops, ordered by a key: which function of `heapq` gives the k smallest items by a key?","The key of a shop is its squared distance, `shop[0] * shop[0] + shop[1] * shop[1]`, and `heapq.nsmallest` keeps shops with equal keys in their order.","Taking the closest shop and removing it, k times over, is O(n·k): a billion steps for these sizes. So is keeping your own sorted list of the k closest so far, when the shops come from the farthest to the closest."],"tests":[{"name":"k_closest([(3, 4), (1, 1), (-2, 0), (0, 5)], 2)","stdin":"","call":"k_closest([(3, 4), (1, 1), (-2, 0), (0, 5)], 2)\n","expect":"[(1, 1), (-2, 0)]","hidden":false,"returns":true},{"name":"k_closest([(3, 4), (1, 1)], 0)","stdin":"","call":"k_closest([(3, 4), (1, 1)], 0)\n","expect":"[]","hidden":true,"returns":true},{"name":"k_closest([(0, 5), (3, 4), (5, 0), (1, 1)], 3)","stdin":"","call":"k_closest([(0, 5), (3, 4), (5, 0), (1, 1)], 3)\n","expect":"[(1, 1), (0, 5), (3, 4)]","hidden":true,"returns":true},{"name":"k_closest([(5, 0), (0, 5), (4, 3), (1, 1)], 3)","stdin":"","call":"k_closest([(5, 0), (0, 5), (4, 3), (1, 1)], 3)\n","expect":"[(1, 1), (5, 0), (0, 5)]","hidden":true,"returns":true},{"name":"k_closest([(2, 2), (-1, 0), (0, -3), (1, -1)], 4)","stdin":"","call":"k_closest([(2, 2), (-1, 0), (0, -3), (1, -1)], 4)\n","expect":"[(-1, 0), (1, -1), (2, 2), (0, -3)]","hidden":true,"returns":true},{"name":"k_closest([(7, 1)], 1)","stdin":"","call":"k_closest([(7, 1)], 1)\n","expect":"[(7, 1)]","hidden":true,"returns":true},{"name":"Test 7","stdin":"","call":"result = k_closest([(i % 1000 - 500, i * 7 % 997 - 498) for i in range(200_000)], 5000)\nprint(len(result), result[:3], result[-1])\n","expect":"5000 [(0, 1), (1, -1), (0, 2)] (5, -89)","hidden":true,"big":true},{"name":"Test 8","stdin":"","call":"result = k_closest([(i, -i // 2) for i in range(100_000, 0, -1)], 5000)\nprint(len(result), result[:3], result[-1])\n","expect":"5000 [(1, -1), (2, -1), (3, -2)] (5000, -2500)","hidden":true,"big":true}],"require":[],"forbid":[]},{"task":"Write `median(values)` with quickselect, without sorting: for an odd number of values, the one in the middle of the sorted order; for an even number, the average of the two in the middle. `values` is never empty, and the function must leave it as it is. It must be fast for 100,000 values, also when many of them are equal. Don't use `sorted`, `sort`, `heapq` or `statistics`.\n\nFor example, this program:\n\n```python\nprint(median([120, 45, 80]))\nprint(median([300, 90, 15, 60]))\n```\n\nprints:\n\n```\n80\n75.0\n```\n","kind":"code","seed":"def median(values):\n    pass\n\n\nif __name__ == \"__main__\":\n    print(median([120, 45, 80]))\n    print(median([300, 90, 15, 60]))\n","solution":"def partition3(values, low, high, pivot):\n    lt = low\n    i = low\n    gt = high\n    while i <= gt:\n        if values[i] < pivot:\n            values[lt], values[i] = values[i], values[lt]\n            lt += 1\n            i += 1\n        elif values[i] > pivot:\n            values[i], values[gt] = values[gt], values[i]\n            gt -= 1\n        else:\n            i += 1\n    return lt, gt\n\n\ndef kth_smallest(values, k):\n    values = values[:]\n    low = 0\n    high = len(values) - 1\n    while True:\n        pivot = values[(low + high) // 2]\n        lt, gt = partition3(values, low, high, pivot)\n        if k < lt:\n            high = lt - 1\n        elif k > gt:\n            low = gt + 1\n        else:\n            return pivot\n\n\ndef median(values):\n    n = len(values)\n    if n % 2 == 1:\n        return kth_smallest(values, n // 2)\n    return (kth_smallest(values, n // 2 - 1) + kth_smallest(values, n // 2)) / 2\n\n\nif __name__ == \"__main__\":\n    print(median([120, 45, 80]))\n    print(median([300, 90, 15, 60]))\n","hints":["Write the lesson's `partition3` and `kth_smallest` first: `kth_smallest(values, k)` gives the value at index k of the sorted order.","With `n = len(values)`: for an odd n the median is `kth_smallest(values, n // 2)`; for an even n, add `kth_smallest(values, n // 2 - 1)` and `kth_smallest(values, n // 2)` and divide by 2.","`kth_smallest` must work on a copy, `values[:]`, or the caller's list changes. Use the three-way partition. With the two-way one, a long run of equal values loses only one item per round: for 100,000 values that are all the same, that's billions of steps. Counting, for each value, how many others are smaller is O(n²), billions of steps too."],"tests":[{"name":"Test 1","stdin":"","call":"print(median([120, 45, 80]))\nprint(median([300, 90, 15, 60]))\n","expect":"80\n75.0","hidden":false},{"name":"median([5])","stdin":"","call":"median([5])\n","expect":"5","hidden":true,"returns":true},{"name":"median([3, 3, 3, 8])","stdin":"","call":"median([3, 3, 3, 8])\n","expect":"3.0","hidden":true,"returns":true},{"name":"median([10, -4])","stdin":"","call":"median([10, -4])\n","expect":"3.0","hidden":true,"returns":true},{"name":"Test 5","stdin":"","call":"values = [9, 2, 7, 4, 5]\nprint(median(values), values)\n","expect":"5 [9, 2, 7, 4, 5]","hidden":true},{"name":"median(list(range(1001)))","stdin":"","call":"median(list(range(1001)))\n","expect":"500","hidden":true,"returns":true},{"name":"median([i * i % 1_000_003 for i in range(100_001)])","stdin":"","call":"median([i * i % 1_000_003 for i in range(100_001)])\n","expect":"497705","hidden":true,"big":true,"returns":true},{"name":"median([i * i % 1_000_003 for i in range(100_000)])","stdin":"","call":"median([i * i % 1_000_003 for i in range(100_000)])\n","expect":"497696.5","hidden":true,"big":true,"returns":true},{"name":"median([i % 3 for i in range(100_000)])","stdin":"","call":"median([i % 3 for i in range(100_000)])\n","expect":"1.0","hidden":true,"big":true,"returns":true},{"name":"median([7] * 100_000)","stdin":"","call":"median([7] * 100_000)\n","expect":"7.0","hidden":true,"big":true,"returns":true}],"require":[],"forbid":[{"pattern":"\\bsorted\\s*\\(|\\.sort\\s*\\(","message":"Find the median with quickselect, without `sorted` or `sort`."},{"pattern":"\\bheapq\\b|\\bstatistics\\b","message":"Write quickselect yourself on this step, without `heapq` or `statistics`."}]}]}