Optimizing Binary Search: Achieving Mechanical Sympathy

By Itamar Turner-Trauring
Updated: July 15, 2026 | Originally published: July 11, 2026

When tasked with accelerating computational Python, the most intuitive first step is usually: Select a superior algorithm.

The Real-World Challenge

Consider a specific bottleneck found within scikit-learn’s gradient histogram boosting algorithm. The goal is to take a massive array of floating-point numbers and map them to an integer range of $0$ to $254$, ensuring the distribution is even.

The workflow looks like this:

While scikit-learn's version is already written in a compiled language and utilizes multi-core parallelism, there was still room for improvement. Inspired by research from Paul Khuong and my work at Quansight, I managed to significantly boost the performance of this process.

In this guide, I will demonstrate these optimizations using a simplified example, eventually reaching a version that is 6× faster than the original.

Warning: This article moves quickly through low-level hardware concepts. We will touch upon Instruction-Level Parallelism (ILP), SIMD, memory caches, and branch misprediction. This is a demonstration of possibility, not an exhaustive textbook.

---

The Baseline: Standard Binary Search

Although the original scikit-learn code used Cython, I will use Rust for these examples. Below is a classic binary search implementation tailored for bucketization.

use std::cmp::Ordering;

/// Since Rust prevents direct float comparison (due to NaN inconsistencies),
/// we use a helper function for total ordering.
fn less_than(a: f64, b: f64) -> bool {
    a.total_cmp(&b) == Ordering::Less
}

/// Maps floating point values to integer buckets based on provided boundaries.
fn bucketize_classic_impl(arr: &[f64], boundaries: &[f64]) -> Vec<usize> {
    // Pre-allocate memory to avoid multiple re-allocations
    let mut result = Vec::with_capacity(arr.len());

    for value in arr {
        let mut min_idx = 0;
        let mut max_idx = boundaries.len();

        while min_idx < max_idx {
            // Calculate midpoint to avoid potential overflow
            let middle = min_idx + ((max_idx - min_idx) / 2);
            
            if less_than(boundaries[middle], value) {
                min_idx = middle + 1;
            } else {
                max_idx = middle;
            }
        }
        result.push(min_idx);
    }
    result
}

Integrating with Python

To make this usable in a Python environment, we use PyO3 and numpy. While this is mostly boilerplate, it allows the Rust function to accept and return NumPy arrays.

<details> <summary>Click to view the PyO3 Boilerplate</summary>

use pyo3::prelude::*;
use numpy::{ndarray::Array, PyArray1, PyReadonlyArray1};

#[pyfunction]
fn bucketize_classic<'py>(
    py: Python<'py>,
    arr: PyReadonlyArray1<f64>,
    boundaries: PyReadonlyArray1<f64>,
) -> PyResult<Bound<'py, PyArray1<usize>>> {
    let result = bucketize_classic_impl(
        arr.as_slice().unwrap(), 
        boundaries.as_slice().unwrap()
    );
    let result = PyArray1::from_owned_array(py, Array::from_vec(result));
    Ok(result)
}

</details>

---

The Bottleneck: Branch Mispredictions

If we are already using a compiled language and an $O(\log n)$ algorithm, how do we get more speed? The answer lies in mechanical sympathy—designing software that harmonizes with the underlying hardware.

Understanding CPU Execution

Many Python developers view code execution as a linear sequence: One instruction $\rightarrow$ One result $\rightarrow$ Next instruction.

In compiled languages, this is a myth. Modern CPU cores utilize Instruction-Level Parallelism (ILP), allowing them to execute multiple independent instructions simultaneously.

fn two_adds(a: i64, b: i64, c: i64, d: i64) -> i64 {
    // The CPU can likely compute these two sums at the same time
    let t1 = a + b;
    let t2 = c + d;
    t1 + t2
}

The Problem with Branches

The efficiency of ILP collapses when the CPU encounters a branch (if, while, for). The CPU doesn't know which path the code will take until the condition is evaluated, but it doesn't want to stop working.

To solve this, the CPU uses a Branch Predictor.

Consider this example:

fn maybe_add(a: i64, b: i64, c: i64, d: i64, add: bool) -> i64 {
    let t1 = a + b;
    // The CPU must guess whether to run the addition or multiplication
    let t2 = if add { c + d } else { c * d };
    t1 + t2
}

Why Binary Search is Slow

In a standard binary search, the decision to move the min_idx or max_idx is essentially a coin flip (50% probability) if the data is distributed uniformly.

The branch predictor cannot find a pattern in these flips. Consequently, the CPU frequently guesses wrong, flushes its pipeline, and restarts, wasting precious clock cycles.

Optimization Roadmap:

• Implement basic binary search
• Eliminate unpredictable branches
• Optimize memory access patterns
• Leverage SIMD instructions