The k lightest crates

An orchard packhouse weighs every crate coming off the grader. The small van takes the three lightest, and the run holds a million crates.

The problem

Crates leave the grader in no particular weight order — the grader sorts by size, and weight only correlates loosely with it. The dispatcher has a van going to a single small shop and wants the k lightest crates in the run, listed by increasing weight, so the loaders can pick them off in one walk down the line.

Two crates can weigh the same, and both count: the van takes crates, not distinct weights. k may be zero, in which case nothing is loaded, and k may exceed the number of crates, in which case the van takes the whole run and the list is every weight in order.

Input. crates — a list of integers, the weight in kilograms of each crate in grader order. k — how many crates the van takes.

Output. A list of the min(k, len(crates)) smallest weights, in increasing order.

Example.

crates = [12, 5, 9, 5, 20, 3], k = 3   ->  [3, 5, 5]

Both 5 kg crates are loaded; there is no de-duplication.

A second example, at the two ends of the range for k:

crates = [12, 5, 9, 5, 20, 3], k = 9   ->  [3, 5, 5, 9, 12, 20]
crates = [12, 5, 9, 5, 20, 3], k = 0   ->  []

Constraints.

• 0 <= len(crates) <= 10^6
• 1 <= crates[i] <= 10^4
• 0 <= k <= 10^6

Hints

Approach

Brute force

Sort the run and slice the first k. That is about n log n comparisons — some 2·10⁷ for a million crates — and if k is 3, then 999 997 of the positions you paid to order are discarded unread.

The insight

Only the first k positions need to be ordered, and a heap over the whole run can be built in O(n) rather than O(n log n), so the total is O(n + k log n).

The linear build works because the tree is bottom-heavy: half the crates are leaves and sift down zero levels, a quarter sift down at most one, an eighth at most two, and the sum of those levels converges to n. It needs no precondition at all — any arrangement of the run can be heapified, which is exactly why the build is cheaper than inserting the crates one by one.

Algorithm

1. If k is zero or the run is empty, return the empty list.
2. Copy the weights, because heapify rearranges its argument in place.
3. Heapify the copy in one pass.
4. Pop min(k, n) times, collecting each popped weight.
5. The pops already come out in increasing order.

Complexity

Time O(n + k log n) — a linear build plus k pops of log n each; for k = 3 on a million crates that is a million steps rather than twenty million. Space O(n) for the copy, or O(1) extra if rearranging the caller's list is allowed.

Solution

"""The k lightest crates — heapify the whole run once, then pop k times."""

import heapq


def solve(crates, k):
    if k <= 0:
        return []
    heap = list(crates)      # copy: heapify rearranges its argument in place
    heapq.heapify(heap)      # O(n): half the nodes are leaves and sift down zero levels
    k = min(k, len(heap))
    # Invariant: after each pop the root is the lightest crate not yet listed,
    # so the crates come out already in increasing order.
    return [heapq.heappop(heap) for _ in range(k)]

TESTS = [
    (([12, 5, 9, 5, 20, 3], 3), [3, 5, 5]),
    (([12, 5, 9, 5, 20, 3], 6), [3, 5, 5, 9, 12, 20]),
    (([12, 5, 9, 5, 20, 3], 9), [3, 5, 5, 9, 12, 20]),
    (([12, 5, 9, 5, 20, 3], 0), []),
    (([4, 4, 4, 4], 2), [4, 4]),
    (([7], 1), [7]),
    (([], 3), []),
]

Pitfalls

• Heapifying the caller's list. heapq.heapify rearranges in place, so the dispatcher's grader-order record is destroyed and the next report reads a shuffled run. Copy first.
• Reading heap[:k] instead of popping. A heap is not sorted. After heapifying the first example the array is [3, 5, 9, 5, 20, 12], so the first three entries give [3, 5, 9] — and 9 kg is not among the three lightest.
• Popping k times without clamping. With six crates and k = 9, the seventh pop raises IndexError on an empty heap.

Variants

• The nearest drill bits — top-k again, but the heap is capped at k and keyed by a computed distance rather than by the value itself.
• Top-k and two heaps — when the size-k heap beats heapifying everything, costed both ways.