efficient way to get all numpy slices for different ranges

user_7812697 0
Viewed 135

I want to slice the same numpy array (data_arra) multiple times to find each time the values in a different range

data_ar shpe: (203,)

range_ar shape: (1000,)

I implemented it with a for loop, but it takes way to long since I have a lot of data_arrays:

#create results array
results_ar = np.zeros(shape=(1000),dtype=object)

i=0
for range in range_ar:
    results_ar[i] = data_ar[( (data_ar>=(range-delta)) & (data_ar<(range+delta)) )].values
    i+=1

so for example:

data_ar = [1,3,4,6,10,12]
range_ar = [7,4,2]
delta= 3

expected output:
(note results_ar shpae=(3,) dtype=object, each element is an array)

results_ar[[6,10];
           [1,3,4,6];
           [1,3,4]]

some idea on how to tackle this?

numpy-slicing numpy python
2 Answers

You can use numba to speed up the computations.

import numpy as np
import numba
from numba.typed import List
import timeit

data_ar = np.array([1,3,4,6,10,12])
range_ar = np.array([7,4,2])
delta = 3

def foo(data_ar, range_ar):
    results_ar = list()
    for i in range_ar:
        results_ar.append(data_ar[( (data_ar>=(i-delta)) & (data_ar<(i+delta)) )])

print(timeit.timeit(lambda :foo(data_ar, range_ar)))

@numba.njit(parallel=True, fastmath=True)
def foo(data_ar, range_ar):
    results_ar = List()
    for i in range_ar:
        results_ar.append(data_ar[( (data_ar>=(i-delta)) & (data_ar<(i+delta)) )])

print(timeit.timeit(lambda :foo(data_ar, range_ar)))

15.53519330600102

1.6557575029946747

An almost 9.8 times speedup.

user_88588320

You could use np.searchsorted like this:

data_ar = np.array([1, 3, 4, 6, 10, 12])
range_ar = np.array([7, 4, 2])
delta = 3

bounds = range_ar[:, None] + delta * np.array([-1, 1])

result = [data_ar[slice(*row)] for row in np.searchsorted(data_ar, bounds)]
user_29887300
Related