Trouver la distance au zéro le plus proche dans le tableau NumPy

Disons que j'ai un tableau NumPy:

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

À chaque index, je veux trouver la distance à la valeur zéro la plus proche. Si la position est un zéro lui-même, retournez zéro comme distance. Par la suite, nous nous intéressons uniquement aux distances au zéro le plus proche situé à droite de la position actuelle. L'approche super naïve serait quelque chose comme:

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

Et la sortie serait:

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

Je remarque un schéma de compte à rebours / décrément dans la sortie entre les zéros. Donc, je pourrais peut-être utiliser les emplacements des zéros (c.-à-d. zero_indices = np.argwhere(x == 0).flatten())

Quel est le moyen le plus rapide d'obtenir la sortie souhaitée en temps linéaire?

Approche n ° 1: Searchsorted à la rescousse du temps linéaire de manière vectorisée (avant que les gars de Numba n'entrent)!

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

Approche n ° 2: une autre avec certains 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))

Alternativement, la dernière étape de cumsumpourrait être remplacée par des repeatfonctionnalités -

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

Approche n ° 3: une autre avec surtout juste 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

Vous pourriez travailler de l'autre côté. Gardez un compteur sur le nombre de chiffres non nuls passés et affectez-le à l'élément du tableau. Si vous voyez 0, remettez le compteur à 0

Edit: s'il n'y a pas de zéro à droite, alors vous avez besoin d'un autre contrôle

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)

Vous pouvez utiliser la différence entre les indices de chaque position et le maximum cumulé des positions nulles pour déterminer la distance au zéro précédent. Cela peut être fait en avant et en arrière. La distance minimale entre l'avant et l'arrière au zéro précédent (ou suivant) sera la plus proche:

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)

résultats:

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]

Cas particulier où aucun zéro n'est présent sur les bords extérieurs:

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]

fonctionne également sans aucun zéro

[EDIT]  solutions non numpy ...

si vous recherchez une solution O (N) qui ne nécessite pas numpy, vous pouvez appliquer cette stratégie en utilisant la fonction d'accumulation d'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)

production:

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 vous ne souhaitez utiliser aucune bibliothèque, vous pouvez accumuler les distances manuellement dans une boucle:

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)

production:

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]

Ma première intuition serait d'utiliser le tranchage. Si x peut être une liste normale au lieu d'un tableau numpy, alors vous pouvez utiliser

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

si numpy est nécessaire, vous pouvez utiliser

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

mais cela est moins efficace car vous recherchez tous les emplacements zéro à droite de la valeur, puis vous retirez uniquement le premier. Certainement une meilleure façon de le faire en numpy.

Edit: je suis désolé, j'ai mal compris. Cela vous donnera la distance aux zéros les plus proches - que ce soit à gauche ou à droite. Mais vous pouvez utiliser d_rightcomme résultat intermédiaire. Cela ne couvre pas le cas de bord de ne pas avoir de zéro à droite.

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])