6× faster binary search: from compiled code to mechanical sympathy
6× faster binary search: from compiled code to mechanical sympathy
by Itamar Turner-Trauring<br>Last updated 11 Jul 2026, originally created 11 Jul 2026
How do you speed up computational Python code? A common, and useful,<br>starting point is:
Pick a good algorithm.
Use a compiled language to write a Python extension.
Maybe add parallelism so you can use multiple CPU cores.
But what if you need more speed? Consider the following real problem,<br>one of the steps in scikit-learn’s gradient histogram boosting<br>algorithm:
You have a large array of floating point numbers.
You want to assign them to the integer range 0-254, spread out evenly.
scikit-learn implements this by splitting up the full range of float<br>values into 255 buckets, creating a sorted array of bucket boundaries,<br>and then using binary search to choose the appropriate bucket for each<br>value. The binary search is implemented in a compiled language, and it<br>can run in parallel on multiple cores.
Recently, as part of my work at<br>Quansight, and inspired by<br>two<br>posts<br>by Paul Khuong, I sped up this implementation significantly. How? By<br>making sure the code wasn’t fighting against the CPU.
In this article I’m going to walk you through that speed-up,<br>demonstrated on a simplified example. Then I’m going to demonstrate a<br>series of additional optimizations, with the final version running 6×<br>faster than the original one.
It’s worth knowing that I will be speeding through mentions of many<br>different low-level hardware topics: instruction-level parallelism,<br>branch (mis)prediction, memory caches, SIMD, and more. This is only one<br>article, it can only briefly introduce you to what’s possible, it can’t<br>function as an in-depth tutorial. So I’ll talk about how you can learn<br>more about these topics at the end of the article.
The starting point: Standard binary search
The original scikit-learn code was implemented in Cython, but I’m going<br>to use Rust for this article. Here’s a pretty standard implementation of<br>binary search (based on the one in NumPy), designed for this use case of<br>finding a bucket given an array of boundaries:
use std::cmp::Ordering;
/// Rust doesn't let you compare floats with the normal<br>/// operator (because NaN makes comparison results<br>/// inconsistent), so implement a custom function to do so.<br>fn less_than(a: f64, b: &f64) -> bool {<br>a.total_cmp(b) == Ordering::Less
/// Convert floating points values into integer values by<br>/// finding which bucket they fit in, given bucket<br>/// boundaries.<br>fn bucketize_classic_impl(<br>arr: &[f64],<br>boundaries: &[f64],<br>) -> Vecusize> {<br>// A Vec or vector is Rust's equivalent of a Python<br>// list. Here I create an empty Vec with enough memory<br>// allocated to store `arr.len()` values:<br>let mut result = Vec::with_capacity(arr.len());<br>for value in arr {<br>// Standard binary search algorithm:<br>let mut min_idx = 0;<br>let mut max_idx = boundaries.len();<br>while min_idx max_idx {<br>let middle = min_idx + ((max_idx - min_idx) / 2);<br>if less_than(boundaries[middle], value) {<br>min_idx = middle + 1;<br>} else {<br>max_idx = middle;<br>// This is equivalent to a_list.append() in Python:<br>result.push(min_idx);<br>// Return the result:<br>result
For completeness, here’s how you can hook it up to Python, so it takes<br>and returns NumPy arrays; this is boilerplate, so I’m not going to show<br>it for later functions. You can just skip it if you’re not familiar with<br>Rust and PyO3, or don’t particularly care; it’s not relevant to the rest<br>of the article.
Click here to see the code
use pyo3::prelude::*;<br>use numpy::ndarray::Array;<br>use numpy::{PyArray1, PyReadonlyArray1};
#[pyfunction]<br>fn bucketize_classic'py>(<br>py: Python'py>,<br>// These two arguments are 1-dimensional arrays of<br>// floats:<br>arr: PyReadonlyArray1f64>,<br>boundaries: PyReadonlyArray1f64>,<br>) -> PyResultBound'py, PyArray1usize>>> {<br>let result = bucketize_classic_impl(<br>arr.as_slice().unwrap(),<br>boundaries.as_slice().unwrap(),<br>);<br>// Convert a Rust vector into a 1D NumPy array we can<br>// pass back to Python:<br>let result =<br>PyArray1::from_owned_array(py, Array::from_vec(result));<br>Ok(result)
Branch mispredictions slow your code down
How can I speed up this implementation of binary search? It’s already<br>using a scalable algorithm, and a compiled language. Parallelism is<br>certainly an option, but I’m going to use a different approach:<br>mechanical sympathy, a better understanding of how the CPU works. I’ll<br>start with a very quick review of how modern CPUs run code in parallel<br>within a single core.
A reasonable mental model of Python code is that the code is executed<br>one instruction at a time. Do twice as many arithmetic operations, and<br>the code will run twice as slow. Once you switch to a compiled language,<br>with operations sometimes mapping to one or two CPU instructions, that<br>mental model is no longer correct. Modern CPUs can sometimes run<br>multiple independent CPU instructions at once (“instruction-level<br>parallelism”) on a single CPU core, resulting in...