Skip to main content
 首页 » 编程设计

python之在numpy数组的每一行中获取第一个零的索引数组

2025年05月04日26soundcode

我有一个包含 0 和 255 的 1650 行和 1275 列的 numpy 数组。 我想获取行中每个第一个零的索引并将其存储在数组中。 我使用 for 循环来实现这一点。这是示例代码

#new_arr is a numpy array and k is an empty array  
for i in range(new_arr.shape[0]): 
  if not np.all(new_arr[i,:]) == 255: 
   x = np.where(new_arr[i,:]==0)[0][0] 
   k.append(x) 
  else: 
   k.append(-1) 

1650 行大约需要 1.3 秒。有没有其他方法或函数可以更快地获取索引数组?

请您参考如下方法:

一种方法是使用 ==0 获取匹配掩码,然后获取 argmax沿着每一行,即 argmax(axis=1) 为我们提供每一行的第一个匹配索引 -

(arr==0).argmax(axis=1) 

sample 运行-

In [443]: arr 
Out[443]:  
array([[0, 1, 0, 2, 2, 1, 2, 2], 
       [1, 1, 2, 2, 2, 1, 0, 1], 
       [2, 1, 0, 1, 0, 0, 2, 0], 
       [2, 2, 1, 0, 1, 2, 1, 0]]) 
 
In [444]: (arr==0).argmax(axis=1) 
Out[444]: array([0, 6, 2, 3]) 

捕获非零行(如果可以的话!)

为了方便不包含任何零的行,我们需要再做一步工作,并进行一些屏蔽 -

In [445]: arr[2] = 9 
 
In [446]: arr 
Out[446]:  
array([[0, 1, 0, 2, 2, 1, 2, 2], 
       [1, 1, 2, 2, 2, 1, 0, 1], 
       [9, 9, 9, 9, 9, 9, 9, 9], 
       [2, 2, 1, 0, 1, 2, 1, 0]]) 
 
In [447]: mask = arr==0 
 
In [448]: np.where(mask.any(1), mask.argmax(1), -1) 
Out[448]: array([ 0,  6, -1,  3])