Numpy有序一维数组：无循环筛选A中与B元素差绝对值≤10的元素

针对两个已排序数组的场景，最优的无循环实现可以基于np.searchsorted二分查找能力完成，完全不需要显式遍历，且不受数组长度差、元素位置偏移的影响：

核心实现思路 #

由于B是已升序排序的数组，对A中任意元素a，只要判断B中是否存在元素落在[a-10, a+10]区间即可。利用二分查找可以快速定位区间的左右边界：

• 用np.searchsorted找a-10在B中的左插入位置
• 用np.searchsorted找a+10在B中的右插入位置
• 如果左插入位置小于右插入位置，说明两个位置之间存在至少一个B的元素，即满足abs(a-b) <=10的条件

代码示例 #

import numpy as np

# 示例已排序数组
A = np.array([2, 12, 25, 40, 58])
B = np.array([8, 22, 48])

# 批量计算所有A元素对应的区间边界
left_pos = np.searchsorted(B, A - 10, side='left')
right_pos = np.searchsorted(B, A + 10, side='right')

# 筛选符合条件的A元素
filtered_A = A[left_pos < right_pos]

上述代码输出为array([ 2, 12, 25, 40])，符合预期。

方案优势 #

• 无Python层面的循环，底层为C实现，执行效率很高，时间复杂度为O(len(A) * log len(B))
• 对整数、浮点数类型的数组都适用，不受数值范围限制
• 内存占用极低，不需要构造超大掩码数组

如果你确实想使用掩码数组的思路，仅适用于元素为整数且数值范围不大的场景，实现代码如下：

min_val = min(A.min(), B.min())
max_val = max(A.max(), B.max())
# 构造B的存在性掩码
b_mask = np.zeros(max_val - min_val + 1, dtype=bool)
b_mask[B - min_val] = True
# 生成±10的窗口掩码
window_mask = np.convolve(b_mask, np.ones(21, dtype=bool), mode='same')
# 筛选A中符合条件的元素
filtered_A = A[window_mask[A - min_val]]

该方案在数值范围较大时会占用大量内存，不推荐作为通用方案使用。

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