Find Distance to Nearest Zero in NumPy Array

You can use the difference between the indices of each position and the cumulative max of zero positions to determine the distance to the preceding zero. This can be done forward and backward. The minimum between forward and backward distance to the preceding (or next) zero will be the nearest:

import numpy as np

indices  = np.arange(x.size)
zeroes   = x==0
forward  = indices - np.maximum.accumulate(indices*zeroes)  # forward distance
forward[np.cumsum(zeroes)==0] = x.size-1                    # handle absence of zero from edge
forward  = forward * (x!=0)                                 # set zero positions to zero                

zeroes   = zeroes[::-1]
backward = indices - np.maximum.accumulate(indices*zeroes) # backward distance
backward[np.cumsum(zeroes)==0] = x.size-1                  # handle absence of zero from edge
backward = backward[::-1] * (x!=0)                         # set zero positions to zero

distZero = np.minimum(forward,backward) # closest distance (minimum)

results:

distZero
# [0, 1, 1, 0, 1, 2, 2, 1, 0, 0]

forward
# [0, 1, 2, 0, 1, 2, 3, 4, 0, 0]

backward
# [0, 2, 1, 0, 4, 3, 2, 1, 0, 0]

Special case where no zeroes are present on outer edges:

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

forward:  [9 9 9 0 1 2 3 0 1 2]
backward: [3 2 1 0 3 2 1 0 9 9]
distZero: [3 2 1 0 1 2 1 0 1 2]

also works with no zeroes at all

[EDIT] non-numpy solutions ...

if you're looking for an O(N) solution that doesn't require numpy, you can apply this strategy using the accumulate function from itertools:

x = [0, 1, 2, 0, 4, 5, 6, 7, 0, 0]

from itertools import accumulate

maxDist  = len(x) - 1
zeroes   = [maxDist*(v!=0) for v in x]
forward  = [*accumulate(zeroes,lambda d,v:min(maxDist,(d+1)*(v!=0)))]
backward = accumulate(zeroes[::-1],lambda d,v:min(maxDist,(d+1)*(v!=0)))
backward = [*backward][::-1]
distZero = [min(f,b) for f,b in zip(forward,backward)]                      

print("x",x)
print("f",forward)
print("b",backward)
print("d",distZero)

output:

x [0, 1, 2, 0, 4, 5, 6, 7, 0, 0]
f [0, 1, 2, 0, 1, 2, 3, 4, 0, 0]
b [0, 2, 1, 0, 4, 3, 2, 1, 0, 0]
d [0, 1, 1, 0, 1, 2, 2, 1, 0, 0]

If you don't want to use any library, you can accumulate the distances manually in a loop:

x = [0, 1, 2, 0, 4, 5, 6, 7, 0, 0]
forward,backward = [],[]
fDist = bDist = maxDist = len(x)-1
for f,b in zip(x,reversed(x)):
    fDist = min(maxDist,(fDist+1)*(f!=0))
    forward.append(fDist)
    bDist = min(maxDist,(bDist+1)*(b!=0))
    backward.append(bDist)
backward = backward[::-1]
distZero = [min(f,b) for f,b in zip(forward,backward)]

print("x",x)
print("f",forward)
print("b",backward)
print("d",distZero)

output:

x [0, 1, 2, 0, 4, 5, 6, 7, 0, 0]
f [0, 1, 2, 0, 1, 2, 3, 4, 0, 0]
b [0, 2, 1, 0, 4, 3, 2, 1, 0, 0]
d [0, 1, 1, 0, 1, 2, 2, 1, 0, 0]

You could work from the other side. Keep a counter on how many non zero digits have passed and assign it to the element in the array. If you see 0, reset the counter to 0

Edit: if there is no zero on the right, then you need another check

x = np.array([0, 1, 2, 0, 4, 5, 6, 7, 0, 0])
out = x 
count = 0 
hasZero = False 
for i in range(x.shape[0]-1,-1,-1):
    if out[i] != 0:
        if not hasZero: 
            out[i] = x.shape[0]-1
        else:
            count += 1
            out[i] = count
    else:
        hasZero = True
        count = 0
print(out)

Approach #1 : Searchsorted to the rescue for linear-time in a vectorized manner (before numba guys come in)!

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

Approach #2 : Another with some 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))

Alternatively, last step of cumsum could be replaced by repeat functionality -

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

Approach #3 : Another with mostly just 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

Tags:

Python

Numpy