获取不同范围内所有numpy切片的有效方法

Question

我想对同一个numpy数组(data_arra)进行多次切片，以便每次找到不同范围内的值

data_ar shpe: (203，)

range_ar shape: (1000，)

我用for循环实现了它，但它需要很长时间，因为我有很多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

所以例如:

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

期望输出:
(注意results_ar shpae=(3，) dtype=object，每个元素都是一个数组)

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

关于如何解决这个问题有什么想法吗?

Answer 1

您可以使用numba来加快计算速度。

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

加速将近9.8倍。

Answer 2

您可以这样使用np.searchsorted:

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)]