NumPy 2.5 makes searchsorted up to 25 times faster by running binary searches in lockstep
Victoria Shinder2 min readnp.searchsorted — NumPy's binary search, used underneath histogramming, interval lookups and parts of SciPy and scikit-learn — got a substantial rewrite in NumPy 2.5. A post on the Scientific Python blog walks through how it was done and reports up to a 25× speedup over NumPy 2.4 in the authors' benchmarks.
What was slow. NumPy 2.4 ran one classic binary search per query key. Each step of a binary search depends on the comparison before it, so the CPU waits on one memory read at a time, and on large arrays those reads usually miss the cache. The old code had one optimisation — reusing the previous bounds when the keys arrive sorted — that helps sorted input and slightly hurts random input.
The idea. Run all the searches together. The authors first wrote it in NumPy itself: keep lo and hi as arrays, one entry per key, and advance all of them each step. That pure-Python vectorised version was already an order of magnitude faster than the native 2.4 code for many keys, because independent searches let the CPU keep several memory accesses in flight. Removing the per-key "still active" mask — every search runs exactly ceil(log2(n)) iterations and converged ones simply stop moving — made it up to twice as fast again.
Making it fit in NumPy. Two arrays of state per query costs memory proportional to the number of keys, which the authors judged unacceptable for a core routine. The fix was to make every interval shrink by the same amount each step, so only the left boundary needs storing and the length is a single shared number — the branchless formulation also described in Algorithmica's binary-search case study. In C++ the first iteration is unrolled: every key is compared against the same median, so it is read once, and the result array needs no initialisation pass. Extra memory is O(1) beyond the output. The change went in as PR #30517.
Benchmarks, with their limits. The numbers come from a MacBook Pro with an Apple M1 Pro, 32-bit integer keys and arrays up to 2^30 elements; the best of 50 runs is reported. Against JAX, TensorFlow and PyTorch, the authors call NumPy competitive; the latter two parallelise across threads instead of batching, and all converge once the array outgrows the cache. They suggest an Eytzinger memory layout as the next direction.
The 2.5 line itself received NumPy 2.5.4 on 10 October, a patch release that among other things makes linalg.svd and linalg.lstsq raise on inputs containing infinities instead of returning garbage or hanging.
Sources: Scientific Python blog — https://blog.scientific-python.org/numpy/searchsorted ; NumPy 2.5.4 release notes — https://github.com/numpy/numpy/releases/tag/v2.5.4