Numpy数组布尔运算性能优化：如何提升比较操作速度？

原实现代码：

import numpy as np

a = np.random.random((100000,))
        
def get_first_index(value, arr):
   firstIndex = np.argmax(arr > value)
   if firstIndex <= 0:
      raise Exception('No index found')
   return firstIndex

for i in range(0, 1000):
   get_first_index(0.5, a)

这段代码用于返回数组中首个大于指定值的元素索引。性能分析显示，99%的耗时集中在arr > value的布尔比较操作——因为numpy的向量化比较会遍历整个数组生成布尔数组，即使第一个符合条件的元素出现在数组前部，也会完成全量遍历，造成不必要的开销。

以下是两种更高效的实现方案：

方案1：用Numba加速提前终止的循环 #

通过编写原生循环，找到第一个符合条件的元素就立即返回，避免全量数组遍历。借助Numba将循环编译为机器码，性能可以接近原生C的水平：

import numba
import numpy as np

a = np.random.random((100000,))

@numba.njit
def get_first_index_numba(value, arr):
    for i in range(arr.size):
        if arr[i] > value:
            return i
    raise Exception('No index found')

# 测试调用
for i in range(1000):
    get_first_index_numba(0.5, a)

这种方案的优势在于提前终止遍历，平均时间复杂度远低于原实现的O(n)，尤其适合无序数组且目标元素位置靠前的场景。

方案2：预排序后使用二分查找（适用于可排序场景） #

如果业务允许对数组预先排序，np.searchsorted的二分查找可以将时间复杂度降到O(log n)，是性能最优的选择。注意需要保留原数组的索引：

import numpy as np

a = np.random.random((100000,))
# 预先排序并保存原索引
sorted_indices = np.argsort(a)
sorted_a = a[sorted_indices]

def get_first_index_sorted(value, sorted_arr, sorted_indices):
    idx = np.searchsorted(sorted_arr, value, side='right')
    if idx >= sorted_arr.size:
        raise Exception('No index found')
    return sorted_indices[idx]

# 测试调用
for i in range(1000):
    get_first_index_sorted(0.5, sorted_a, sorted_indices)

此方案仅适用于可以预先排序、且后续查找基于同一排序数组的场景，排序的O(n log n)开销可以被多次查找的收益抵消。

内容的提问来源于stack exchange，提问作者perotom