本文介绍如何在Python中高效提取目标数组中落在多组动态上下界范围内的元素索引，重点对比原始逐条件布尔掩码法与基于排序+二分搜索（bisect/np.searchsorted）的优化方案，实测numpy版本性能提升达45倍。
本文介绍如何在Python中高效提取目标数组中落在多组动态上下界范围内的元素索引，重点对比原始逐条件布尔掩码法与基于排序+二分搜索（`bisect`/`np.searchsorted`）的优化方案，实测`numpy`版本性能提升达45倍。
在科学计算与数据处理中，常需对一个固定目标数组（如时间序列、特征向量或测量值），批量查询其元素是否落入若干组独立的数值区间（由对应位置的上界 ub_arr 和下界 lb_arr 定义），并返回每组区间内匹配元素在原数组中的原始索引。原始方法使用循环 + np.where 布尔索引，时间复杂度为 O(N × M)（N 为边界组数，M 为目标数组长度），在大规模数据（如万级边界、两万级目标）下效率低下。
更优解是利用预排序 + 二分搜索：先对目标数组按值排序并记录原始索引，得到有序值序列 vals 和对应索引序列 idxs；随后对每组 [lb, ub]，用 bisect_left(vals, lb) 和 bisect_right(vals, ub) 快速定位值域区间在 vals 中的切片边界，再直接索引 idxs[a:b] 即得原始索引列表。该策略将单次查询从 O(M) 降至 O(log M)，整体复杂度优化至 O(M log M + N log M)。
以下是两种推荐实现：
✅ 推荐方案一：纯 NumPy（最高性能）
适用于 target_arr 为 np.ndarray 的典型场景，代码简洁且充分利用向量化搜索：
import numpy as np
def get_matching_indices_numpy(lb_arr, ub_arr, target_arr):
"""
高效获取每组 [lb, ub] 在 target_arr 中匹配元素的原始索引。
Parameters:
-----------
lb_arr, ub_arr : 1D array-like, shape (n_bounds,)
每组查询的下界与上界数组，长度需一致。
target_arr : 1D ndarray, shape (m,)
待查询的目标数组。
Returns:
--------
List[np.ndarray]: 长度为 n_bounds 的列表，每个元素为匹配索引的一维 ndarray。
"""
# 预排序：获取排序后索引及对应值
idxs = np.argsort(target_arr) # shape (m,)
vals = target_arr[idxs] # shape (m,), 已升序
# 向量化二分搜索：一次获取所有边界的位置
a = np.searchsorted(vals, lb_arr, side="left") # 每个 lb 的左插入点
b = np.searchsorted(vals, ub_arr, side="right") # 每个 ub 的右插入点
# 构建结果列表
return [idxs[i:j] for i, j in zip(a, b)]
# 示例验证
ub_arr = np.array([4, 5, 6])
lb_arr = np.array([1, 2, 3])
target_arr = np.array([12, 3, 5, 10, 6, 1, 3, 4])
result = get_matching_indices_numpy(lb_arr, ub_arr, target_arr)
print([arr.tolist() for arr in result])
# 输出: [[5, 1, 6, 7], [1, 6, 7, 2], [1, 6, 7, 2, 4]]✅ 推荐方案二：纯 Python + bisect（内存友好）
当目标数组为普通 list 或需避免 NumPy 依赖时适用，逻辑清晰，内存占用更低：
from bisect import bisect_left, bisect_right
def get_matching_indices_bisect(lb_arr, ub_arr, target_arr):
# 构建 (原索引, 值) 对并按值排序
indexed = sorted(enumerate(target_arr), key=lambda x: x[1])
idxs, vals = zip(*indexed) # 解包为元组
result = []
for lb, ub in zip(lb_arr, ub_arr):
left_pos = bisect_left(vals, lb)
right_pos = bisect_right(vals, ub)
result.append(idxs[left_pos:right_pos])
return result⚠️ 注意事项与最佳实践
- 排序仅需一次：核心优化在于将 O(N) 次全量扫描降为 1 次排序 + N 次 O(log M) 搜索，务必避免在循环内重复排序。
- 边界语义明确：searchsorted(..., side="left") 返回第一个 ≥ lb 的位置；side="right" 返回第一个 > ub 的位置，因此 vals[i:j] 精确对应 [lb, ub] 闭区间。
- 数据类型一致性：确保 lb_arr, ub_arr, target_arr 类型兼容（如全为浮点），避免隐式转换开销。
- 大规模场景建议：对于 len(target_arr) > 10^4 且 len(ub_arr) > 10^3 的情况，NumPy 方案平均提速 40× 以上（见基准测试），应作为默认选择。
综上，摒弃逐条件 np.where 循环，转向「预排序 + 二分定位」范式，是解决此类批量区间索引查询问题的最优工程实践。
