NumPy 2.5 : searchsorted jusqu'à 25× plus rapide

Ce qui se passe

np.searchsorted est l'implémentation de la recherche binaire dans NumPy : elle trouve la position d'insertion de chaque clé dans un tableau trié. Les histogrammes et les recherches d'intervalles s'en servent, donc SciPy et scikit-learn en profitent aussi. Les auteurs du correctif racontent comment ils l'ont rendue jusqu'à 25 fois plus rapide dans NumPy 2.5 (pull request #30517).

Pourquoi la version 2.4 attendait

La version 2.4 lance une recherche binaire complète par clé, l'une après l'autre. Dans une recherche, chaque étape dépend de la précédente : le milieu testé décide de la moitié suivante, donc de la prochaine case à lire. Sur un grand tableau, cette lecture sort du cache du processeur, et il attend la mémoire avant de pouvoir continuer.

Le mécanisme

Faire avancer toutes les clés d'un pas à la fois. Les lectures d'un même pas ne dépendent plus les unes des autres, et le processeur en garde plusieurs en vol. L'équipe a d'abord écrit cela en Python vectorisé, qui bat déjà le code natif de la 2.4 d'un ordre de grandeur sur plusieurs clés.

Pour porter la méthode en C++ sans mémoire supplémentaire, elle a réduit l'état : une position de départ par clé, et une longueur d'intervalle commune à toutes les clés, puisque chacune réduit le sien de la même quantité à chaque pas.

import numpy as np

def recherche_groupee(tableau, cles):
    # tableau non vide, trié ; même résultat que np.searchsorted(side="left")
    pos = np.zeros(len(cles), dtype=np.intp)
    reste = len(tableau)              # taille d'intervalle, commune à toutes les clés
    while reste > 1:
        moitie = reste // 2
        milieu = pos + moitie
        pos = np.where(cles > tableau[milieu], milieu, pos)
        reste -= moitie
    return np.where(cles > tableau[pos], pos + 1, pos)

Le port C++ déroule aussi la première itération : toutes les clés sont comparées à la même médiane, lue une fois au lieu de K fois.

Ce que la mesure dit, et ne dit pas

• 25× est un maximum, pas une moyenne.
• Par rapport à la version vectorisée en Python, le C++ va jusqu'à 2× plus vite sur de petits tableaux ; l'écart se réduit quand le tableau grandit.
• Contre PyTorch et TensorFlow, qui répartissent les clés sur des threads (8 cœurs, 20 000 clés dans le test), NumPy 2.5 est compétitif. Une fois le tableau plus grand que le cache, toutes les implémentations se ressemblent.
• Le cas des clés triées n'est pas mesuré. Le code de la 2.4 réutilisait les bornes de la clé précédente quand les clés arrivaient triées ; le code 2.5 montré dans l'article n'en garde pas trace. Si c'est votre usage, chronométrez.

Comment s'y prendre

1. Passer à NumPy 2.5 (notes de version).
2. Chronométrer vos appels à np.searchsorted avant et après, sur vos vraies données : taille du tableau, nombre de clés, clés triées ou non.
3. Si vous avez écrit une recherche binaire à la main, essayer la forme groupée : elle bat le code natif de la 2.4 sans une ligne de C.

Source : Making np.searchsorted up to 25× Faster in NumPy 2.5, blog Scientific Python, octobre 2026.