CODEWITHSUNDEEP
pythonnumpy
roygbiv
roygbiv

Reputation: 172

Find index in sorted array, np.searchsorted() vs. np.where()

I'm trying to find the fastest way to return the index of a given element in a sorted array. np.searchsorted() seems to be specifically designed for the job. It uses binary search, which should be faster than the generic np.where(), which I'm assuming needs to loop through the entire array in O(N). However to my surprise, np.searchsorted() takes the same amount of time as np.where(), even for large arrays. Here's a MWE:

t_where, t_searchsorted = 0, 0

a = sorted(np.random.uniform(0, 1, 10000))  # sorted array

for i in range(1000):
    
    val = np.random.choice(a)  # randomly select value from array

    t = time.time()
    idx_1 = np.where(a == val)[0][0]
    t_where += (time.time()-t)

    t = time.time()
    idx_2 = np.searchsorted(a, val)
    t_searchsorted += (time.time()-t)
    
print(t_where)         # 0.804 s
print(t_searchsorted)  # 0.793 s

EDIT: As @Warren Weckesser explained in his answer below, the issue was that sorted() would convert my numpy array into a python list, and np.searchsorted() was converting it back into a numpy array in O(N). For those who might be wondering, the equivalent of np.searchsorted() for python lists is bisect.bisect_left(). Here's a quick benchmark (reproducing the experiment above):

If a is a (sorted) python list, and val is in the list:

idx = a.index(val)                 # takes 100.0 ms
idx = bisect.bisect_left(a, val)   # takes 1.4 ms

If a is a (sorted) numpy array:

idx = np.where(a, val)         # takes 12.1 ms
idx = np.searchsorted(a, val)  # takes 4.5 ms

Upvotes: 2

Views: 2004

Answers (1)

Warren Weckesser
Warren Weckesser

Reputation: 114791

You used the Python builtin function sorted to sort the random values, so a is a Python list, not a NumPy array. Internally, searchsorted will create a NumPy array from that list, which means copying the list, so the actually complexity is O(n) instead of O(log(n)). The same is true for using np.where: the comparison of the Python list to a numpy.float64 value will end up converting the list to a NumPy array first.

If you change the creation of a to something like

a = np.random.uniform(0, 1, 10000)
a.sort()

you'll get times that are much faster, with relative values that reflect the difference between using where and searchsorted.

Here's the modified version of your script, followed by its output.

import time
import numpy as np


t_where, t_searchsorted = 0, 0

a = np.random.uniform(0, 1, 10000)
a.sort()

for i in range(1000):
    
    val = np.random.choice(a)  # randomly select value from array

    t = time.time()
    idx_1 = np.where(a == val)[0][0]
    t_where += (time.time()-t)

    t = time.time()
    idx_2 = np.searchsorted(a, val)
    t_searchsorted += (time.time()-t)
    
print(t_where)
print(t_searchsorted)

Output:

0.00879216194152832
0.0038056373596191406

Upvotes: 3

Related Questions