[Numpy-discussion] Optimized np.digitize for equidistant bins

[Martín Chalela]

Hi all! I was wondering if there is a way around to using np.digitize when
dealing with equidistant bins. For example:
bins = np.linspace(0, 1, 20)

The main problem I encountered is that digitize calls np.searchsorted. This
is the correct way, I think, for generic bins, i.e. bins that have
different widths. However, in the special, but not uncommon, case of
equidistant bins, the searchsorted call can be very expensive and
unnecessary. One can perform a simple calculation like the following:

def digitize_eqbins(x, bins):
"""
Return the indices of the bins to which each value in input array belongs.
Assumes equidistant bins.
"""
nbins = len(bins) - 1
digit = (nbins * (x - bins[0]) / (bins[-1] - bins[0])).astype(np.int)
return digit + 1

Is there a better way of computing this for equidistant bins?

Thank you!
Martin.
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <https://mail.python.org/pipermail/numpy-discussion/attachments/20201218/6f76fba7/attachment.html>