Is there a better way to search a sorted list if the other list is sorted too?

user_18335909 0
Viewed 117

In the numpy library, one can pass a list into the numpy.searchsorted function, whereby it searched through a different list one element at a time and returns an array of the same sizes as the indices needed to preserve order. However, it seems to be wasting performance if both lists are sorted. For example:

m=[1,3,5,7,9]
n=[2,4,6,8,10]
numpy.searchsorted(m,n)

would return [1,2,3,4,5] which is the correct answer, but it looks like this would have complexity O(n ln(m)), whereby if one were to simply loop through m, and have some kind of pointer to n, it seems like the complexity is more like O(n+m)? Is there some kind of function in NumPy which does this?

numpy python
3 Answers

You could use sortednp, unfortunately it does not give too much flexibility, In the code snippet below I used its merge tracking indices, but it produces three arrays, four times more memory than necessary is used, but it is faster than searchsorted.

import numpy as np
import sortednp as snp

a = np.cumsum(np.random.rand(1000000))
b = np.cumsum(np.random.rand(1000000))

def snp_search(a,b):
  m, (ib, ia) = snp.merge(b, a, indices=True)
  return ib - np.arange(len(ib))

assert(np.all(snp_search(a,b) == np.searchsorted(a,b)))

np.searchsorted(a, b); #58 ms
snp_search(a,b); # 22ms
user_127503530

AFAIK, this is not possible to do that in linear time only with Numpy without making additional assumptions on the inputs (eg. the integer are small and bounded). An alternative solution is to use Numba to do the merge manually:

import numba as nb

# Note: Numba requires a function signature with well defined array types
@nb.njit('int64[:](int64[::1], int64[::1])')
def search_both_sorted(a, b):
    i, j = 0, 0
    result = np.empty(b.size, np.int64)
    while i < a.size and j < a.size:
        if a[i] < b[j]:
            i += 1
        else:
            result[j] = i
            j += 1
    for k in range(j, b.size):
        result[k] = i
    return result

a, b = np.cumsum(np.random.randint(0, 100, (2, 1000000)).astype(np.int64), axis=1)
result = search_both_sorted(a, b)

A faster implementation consists in using a branch-less approach so to remove the overhead of branch mis-prediction (especially on random/unpredictable inputs) when a and b are about the same size. Additionally, the O(n log m) algorithm can be faster when b is small so using np.searchsorted in that case is very efficient as pointed out by @MichaelSzczesny. Note that the Numba implementation of np.searchsorted can be a bit slower than the one of Numpy so it is better to pick the Numpy implementation. Here is the optimized version:

@nb.njit('int64[:](int64[::1], int64[::1])')
def search_both_sorted_opt_numba(a, b):
    sa, sb = a.size, b.size
    # Choose the best algorithm
    if sb < sa * 0.15:
        # Use a version with branches because `a[i] < b[j]`
        # should be most of the time true.
        i, j = 0, 0
        result = np.empty(b.size, np.int64)
        while i < a.size and j < b.size:
            if a[i] < b[j]:
                i += 1
            else:
                result[j] = i
                j += 1
        for k in range(j, b.size):
            result[k] = i
    else:
        # Use a branchless approach to avoid miss-predictions
        i, j = 0, 0
        result = np.empty(b.size, np.int64)
        while i < a.size and j < b.size:
            tmp = a[i] < b[j]
            result[j] = i
            i += tmp
            j += ~tmp
        for k in range(j, b.size):
            result[k] = i
    return result

def search_both_sorted_opt(a, b):
    sa, sb = a.size, b.size
    # Choose the best algorithm
    if 2 * sb * np.log2(sa) < sa + sb:
        return np.searchsorted(a, b)
    else:
        return search_both_sorted_opt_numba(a, b)

searchsorted:                   19.1 ms
snp_search:                     11.8 ms
search_both_sorted:              6.5 ms
search_both_sorted_branchless:   4.3 ms

The optimized branchless Numba implementation is about 4.4 times faster than searchsorted which is pretty good considering that the code of searchsorted is already highly optimized. It can be even faster when a and b are huge because of cache locality.

user_129395570

np.searchsorted takes this into account already as can be seen from the source code:

/*
 * Updating only one of the indices based on the previous key
 * gives the search a big boost when keys are sorted, but slightly
 * slows down things for purely random ones.
 */
if (cmp(last_key_val, key_val)) {
    max_idx = arr_len;
}
else {
    min_idx = 0;
    max_idx = (max_idx < arr_len) ? (max_idx + 1) : arr_len;
}

Here min_idx, max_idx are used to perform binary search on the array. If last_key_val < key_val then only max_idx is reset to the array length, but min_idx remains at its current value, i.e. binary search starts at the same lower boundary as for the previous key.

user_37672390
Related