Skip to main content
 首页 » 编程设计

python之使用 numpy.where 向量化 for 循环

2025年05月04日23davidwang456

>>> a = np.array([1, 1, 2, 3, 3, 4, 5]) 
>>> b = np.array([1, 3, 4, 5]) 
>>> indices = np.zeros((len(a))) 
>>> for i in range(len(a)): 
    try: 
        indices[i] = np.where(b == a[i])[0][0] 
    except: 
        indices[i] = -1 
 
>>> indices 
array([ 0.,  0., -1.,  1.,  1.,  2.,  3.]) 

对于 np.array 的每个元素,我怎样才能在另一个 np.array 中获取它的索引并放置一个常量值(这里,- 1) 后面那个元素哪里找不到?

请您参考如下方法:

使用np.searchsorted获取索引,然后通过将索引索引与 a 进行比较并将它们设置为 -1 -

来重新使用这些索引来获取无效索引
idx = np.searchsorted(b,a) 
idx[b[idx] != a] =-1