我有一个包含 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])
