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

Question

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

data_ar 格式：(203,)

range_ar 形状：(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)]

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

efficient way to get all numpy slices for different ranges

python

numpy

numpy-slicing