Encuentre la distancia al cero más cercano en la matriz NumPy

Digamos que tengo una matriz NumPy:

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

En cada índice, quiero encontrar la distancia al valor cero más cercano. Si la posición es un cero, devuelva cero como distancia. Después, solo nos interesan las distancias al cero más cercano que está a la derecha de la posición actual. El enfoque súper ingenuo sería algo como:

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

Y la salida sería:

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

Estoy notando un patrón de cuenta regresiva / decremento en la salida entre los ceros. Por lo tanto, podría usar las ubicaciones de los ceros (es decir, zero_indices = np.argwhere(x == 0).flatten())

¿Cuál es la forma más rápida de obtener la salida deseada en tiempo lineal?

Enfoque # 1: ¡ Searchsorted al rescate por tiempo lineal de manera vectorizada (antes de que entren los chicos de 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

Enfoque # 2: Otro con algunos 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))

Alternativamente, el último paso de cumsumpodría ser reemplazado por la repeatfuncionalidad:

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

Enfoque # 3: Otro con mayormente solo 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

Podrías trabajar desde el otro lado. Mantenga un contador sobre cuántos dígitos distintos de cero han pasado y asígnelo al elemento en la matriz. Si ve 0, restablezca el contador a 0

Editar: si no hay cero a la derecha, entonces necesita otra verificación

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)

Puede usar la diferencia entre los índices de cada posición y el máximo acumulado de las posiciones cero para determinar la distancia al cero anterior. Esto se puede hacer hacia adelante y hacia atrás. El mínimo entre la distancia hacia adelante y hacia atrás al cero anterior (o siguiente) será el más cercano:

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)

resultados:

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]

Caso especial donde no hay ceros en los bordes exteriores:

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]

también funciona sin ceros en absoluto

[EDITAR]  soluciones no numpy ...

Si está buscando una solución O (N) que no requiera numpy, puede aplicar esta estrategia utilizando la función de acumulación de 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)

salida:

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]

Si no desea utilizar ninguna biblioteca, puede acumular las distancias manualmente en un bucle:

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)

salida:

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]

Mi primera intuición sería usar rebanar. Si x puede ser una lista normal en lugar de una matriz numpy, entonces podría usar

 out = [x[i:].index(0) for i,_ in enumerate(x)]

si es necesario numpy, entonces puede usar

 out = [np.where(x[i:]==0)[0][0] for i,_ in enumerate(x)]

pero esto es menos eficiente porque está buscando todas las ubicaciones cero a la derecha del valor y luego extrae solo la primera. Definitivamente, una mejor manera de hacer esto en numpy.

Editar: Lo siento, no lo entendí. Esto le dará la distancia a los ceros más cercanos, ya sea a la izquierda o a la derecha. Pero puedes usarlo d_rightcomo resultado intermedio. Sin embargo, esto no cubre el caso límite de no tener ningún cero a la derecha.

import numpy as np

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

# Get the distance to the closest zero from the left:
zeros = x == 0
zero_locations = np.argwhere(x == 0).flatten()
zero_distances = np.diff(np.insert(zero_locations, 0, 0))

temp = x.copy()
temp[~zeros] = 1
temp[zeros] = -(zero_distances-1)
d_left = np.cumsum(temp) - 1

# Get the distance to the closest zero from the right:
zeros = x[::-1] == 0
zero_locations = np.argwhere(x[::-1] == 0).flatten()
zero_distances = np.diff(np.insert(zero_locations, 0, 0))

temp = x.copy()
temp[~zeros] = 1
temp[zeros] = -(zero_distances-1)
d_right = np.cumsum(temp) - 1
d_right = d_right[::-1]

# Get the smallest distance from both sides:
smallest_distances = np.min(np.stack([d_left, d_right]), axis=0)
# np.array([0, 1, 1, 0, 1, 2, 2, 1, 0, 0])