The nearest drill bits in the catalogue

A machinist needs a 5.12 mm hole and the supplier does not stock that size. The catalogue is a hundred thousand entries long and in no useful order.

The problem

The supplier's catalogue lists distinct stocked diameters in micrometres, in the order the lines were added to stock, which has nothing to do with size. The machinist enters a required diameter and asks for the k stocked sizes nearest to it.

Nearness is the absolute difference in micrometres. When two sizes are equally far off — one under, one over — the smaller bit is offered, because an undersized hole can be opened out with a reamer and an oversized one cannot be made smaller. The offered sizes are listed in increasing diameter, not in order of nearness, so the machinist can read them off against the rack.

Input. rack — a list of distinct integers, stocked diameters in micrometres. target — the required diameter, which need not be stocked. k — how many sizes to offer.

Output. The k nearest stocked diameters, in increasing order.

Example.

rack = [4800, 5000, 6200, 3400, 5500], target = 5120, k = 3
  ->  [4800, 5000, 5500]

The distances are 320, 120, 1080, 1720 and 380. The three nearest are 5000, 4800 and 5500 — listed by diameter, not by how near they are.

A second example, where the last slot is a tie:

rack = [1800, 2200, 1920, 2080], target = 2000, k = 3   ->  [1800, 1920, 2080]

1920 and 2080 are both 80 off and both go in. The third slot is a tie between 1800 and 2200, each 200 off, and the smaller bit takes it.

Constraints.

• 0 <= len(rack) <= 10^5
• 100 <= rack[i] <= 300000, all distinct
• 100 <= target <= 300000
• 0 <= k <= len(rack)

Hints

Approach

Brute force

Compute every distance, sort the whole catalogue by (distance, diameter), take the first k and sort those by diameter. That is n log n comparisons and n extra pairs in memory — about 1.7 million comparisons and 100 000 pairs to answer a question whose answer is three numbers long.

The insight

Hold exactly k candidates with the worst of them at a root: a new size then costs one comparison, and log k only when it beats that worst.

The key is (distance, diameter) and the heap is ordered with the largest key at the root — a max-heap, which heapq gives by negating both parts of the key. A size that cannot beat the current worst can never enter the answer, because the held set only improves as the read goes on. That is the precondition: an entry's key is fixed and does not depend on what arrives later, so an eviction is final.

Algorithm

1. If k is zero, offer nothing.
2. For each diameter, push (-distance, -diameter) onto a min-heap.
3. If the heap now holds more than k entries, pop: its root is the worst candidate held, so that is the one to drop.
4. After the read, take the diameters back out and sort them ascending.

Complexity

Time O(n log k) for the read, plus O(k log k) for the final sort — with k = 3 the heap is two levels deep, so the read is effectively linear. Space O(k): the catalogue is read once and never copied.

Solution

"""The nearest drill bits — a max-heap capped at k, keyed by distance from the target."""

import heapq


def solve(rack, target, k):
    if k <= 0:
        return []

    # Stored negated, so the ROOT is the worst candidate held: the farthest from
    # target, and among equal distances the larger bit. Invariant: after each
    # entry the heap holds the best min(k, seen) bits, so a bit that loses to the
    # root can never come back.
    worst_first = []
    for diameter in rack:
        heapq.heappush(worst_first, (-abs(diameter - target), -diameter))
        if len(worst_first) > k:
            heapq.heappop(worst_first)

    return sorted(-diameter for _, diameter in worst_first)

TESTS = [
    (([4800, 5000, 6200, 3400, 5500], 5120, 3), [4800, 5000, 5500]),
    (([1800, 2200, 1920, 2080], 2000, 3), [1800, 1920, 2080]),
    (([1800, 2200, 1920, 2080], 2000, 2), [1920, 2080]),
    (([1900, 2100], 2000, 1), [1900]),
    (([4800, 5000, 6200, 3400, 5500], 5120, 0), []),
    (([6350], 100, 1), [6350]),
    (([4800, 5000, 6200, 3400, 5500], 5120, 5), [3400, 4800, 5000, 5500, 6200]),
]

Pitfalls

• Pushing (distance, diameter) and popping when the heap exceeds k. heapq pops the smallest, which is the nearest size, so every pop throws away the best candidate held. The first example comes back [3400, 5500, 6200] — the three furthest.
• Negating the distance but not the diameter. Among equally distant entries the root then holds the smaller bit, so the smaller one is evicted: the second example returns [1920, 2080, 2200] instead of [1800, 1920, 2080].
• Returning the heap array as it stands. A heap is not sorted. On the first example the array reads 5500, 5000, 4800 — the right three sizes, in the wrong order for reading off the rack.

Variants

• The k lightest crates — top-k the other way round: heapify everything and pop k times, which wins when k approaches n.
• The board shows the median — two heaps facing each other, where this keeps one bounded heap.