python - 对于排序数组来说，更快地替代 np.where

Question

给定一个大数组 a沿着每一行排序，是否有比 numpy 的 np.where 更快的替代方案查找 min_v <= a <= max_v 的索引？我想利用数组的排序特性应该能够加快速度。

以下是使用 np.where 进行设置的示例在大数组中查找给定的索引。

import numpy as np

# Initialise an example of an array in which to search
r, c = int(1e2), int(1e6)
a = np.arange(r*c).reshape(r, c)

# Set up search limits
min_v = (r*c/2)-10
max_v = (r*c/2)+10

# Find indices of occurrences
idx = np.where(((a >= min_v) & (a <= max_v)))

Answer 1

当我使用 np.searchsorted 与原始示例中的 1 亿个数字以及不是最新的 NumPy 版本 1.12.1 (无法判断较新的版本)时，它是比 np.where 快不了多少:

>>> import timeit
>>> timeit.timeit('np.where(((a >= min_v) & (a <= max_v)))', number=10, globals=globals())
6.685825735330582
>>> timeit.timeit('np.searchsorted(a.ravel(), [min_v, max_v])', number=10, globals=globals())
5.304438766092062

但是，尽管 searchsorted 的 NumPy 文档说 这个函数使用与内置 python bisect.bisect_left 和 bisect.bisect_right< 相同的算法 函数， 后者要快得多:

>>> import bisect
>>> timeit.timeit('bisect.bisect_left(a.base, min_v), bisect.bisect_right(a.base, max_v)', number=10, globals=globals())
0.002058468759059906

因此，我会使用这个:

idx = np.unravel_index(range(bisect.bisect_left(a.base, min_v),
                             bisect.bisect_right(a.base, max_v)), a.shape)