使用np.where匹配数组元素遇两大问题：阈值报错与性能缓慢求助

问题原因 #

1. DeprecationWarning警告触发原因 #

当执行np.where(A[:,0] == B)时，若A[:,0]与B的数组形状无法满足Numpy的广播规则，元素级比较会失败，从而触发该弃用警告——未来Numpy版本会直接抛出错误而非仅警告。常见场景包括：B是多维数组，或A[:,0]与B的长度差异导致广播逻辑不成立。

2. 运行速度缓慢原因 #

直接使用==配合np.where是逐元素线性扫描，时间复杂度为O(M*N)。当数组规模达到千万级时，计算量呈指数级增长，导致运行效率极低。

修复与提速方案 #

方案一：使用np.isin（简洁高效，适合大多数场景） #

np.isin是Numpy专门用于检查元素是否存在于另一数组的函数，内部基于哈希实现，时间复杂度O(M+N)，同时避免广播兼容性问题。

示例代码：

import numpy as np

# 假设A为二维数组，B为一维数组
match_mask = np.isin(A[:, 0], B)
target_indices = np.where(match_mask)[0]

方案二：排序+np.searchsorted（大规模数据最优解） #

通过先对B排序，再用二分查找定位元素，时间复杂度为O(N log N + M log N)，千万级数据下速度远超其他方法。

示例代码：

import numpy as np

# 对B排序（去重可选，能进一步优化查找效率）
sorted_B = np.sort(np.unique(B))
# 获取A[:,0]在排序后B中的插入位置
insert_pos = np.searchsorted(sorted_B, A[:, 0])
# 生成匹配掩码（避免越界并验证元素相等）
match_mask = (insert_pos < len(sorted_B)) & (sorted_B[insert_pos] == A[:, 0])
target_indices = np.where(match_mask)[0]

方案三：Python集合配合列表推导（适合小数据或混合类型场景） #

若A[:,0]元素为可哈希类型（如整数、字符串），将B转为集合后做成员检查，速度优于逐元素==，但大规模数据下不如Numpy向量化方法高效。

示例代码：

import numpy as np

B_set = set(B)
target_indices = np.array([idx for idx, val in enumerate(A[:,0]) if val in B_set])

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