查找 NumPy 数组中到最近零的距离

假设我有一个 NumPy 数组：

x = np.array([0, 1, 2, 0, 4, 5, 6, 7, 0, 0])

在每个索引处，我想找到到最接近的零值的距离。如果位置本身为零，则返回零作为距离。之后，我们只对到当前位置右侧最接近的零的距离感兴趣。超级天真的方法是这样的：

out = np.full(x.shape[0], x.shape[0]-1)
for i in range(x.shape[0]):
    j = 0
    while i + j < x.shape[0]:
        if x[i+j] == 0:
            break
        j += 1
    out[i] = j

输出将是：

array([0, 2, 1, 0, 4, 3, 2, 1, 0, 0])

我注意到零之间的输出中有倒计时/递减模式。所以，我也许可以使用零的位置（即，zero_indices = np.argwhere(x == 0).flatten())

在线性时间内获得所需输出的最快方法是什么？

方法#1： Searchsorted https://docs.scipy.org/doc/numpy/reference/generated/numpy.searchsorted.html以矢量化方式拯救线性时间（在 numba 人员进来之前）！

mask_z = x==0
idx_z = np.flatnonzero(mask_z)
idx_nz = np.flatnonzero(~mask_z)

# Cover for the case when there's no 0 left to the right
# (for same results as with posted loop-based solution)
if x[-1]!=0:
    idx_z = np.r_[idx_z,len(x)]

out = np.zeros(len(x), dtype=int)
idx = np.searchsorted(idx_z, idx_nz)
out[~mask_z] = idx_z[idx] - idx_nz

方法#2：另一个有一些cumsum -

mask_z = x==0
idx_z = np.flatnonzero(mask_z)

# Cover for the case when there's no 0 left to the right
if x[-1]!=0:
    idx_z = np.r_[idx_z,len(x)]

out = idx_z[np.r_[False,mask_z[:-1]].cumsum()] - np.arange(len(x))

或者，最后一步cumsum可以替换为repeat功能-

r = np.r_[idx_z[0]+1,np.diff(idx_z)]
out = np.repeat(idx_z,r)[:len(x)] - np.arange(len(x))

方法#3：另一个大部分只是cumsum -

mask_z = x==0
idx_z = np.flatnonzero(mask_z)

pp = np.full(len(x), -1)
pp[idx_z[:-1]] = np.diff(idx_z) - 1
if idx_z[0]==0:
    pp[0] = idx_z[1]
else:
    pp[0] = idx_z[0]
out = pp.cumsum()

# Handle boundary case and assigns 0s at original 0s places
out[idx_z[-1]:] = np.arange(len(x)-idx_z[-1],0,-1)
out[mask_z] = 0