Znajdź odległość do najbliższego zera w tablicy NumPy

Powiedzmy, że mam tablicę NumPy:

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

Przy każdym indeksie chcę znaleźć odległość do najbliższej wartości zerowej. Jeśli sama pozycja jest zerem, zwróć zero jako odległość. Potem interesują nas tylko odległości do najbliższego zera, który znajduje się na prawo od aktualnej pozycji. Super naiwne podejście wyglądałoby tak:

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

I wynik byłby:

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

Zauważam wzorzec odliczania / zmniejszania w danych wyjściowych pomiędzy zerami. Mogę być w stanie użyć lokalizacji zer (tj. zero_indices = np.argwhere(x == 0).flatten())

Jaki jest najszybszy sposób na uzyskanie pożądanej wydajności w czasie liniowym?

Podejście nr 1: Searchsorted do ratowania czasu liniowego w wektorowy sposób (zanim pojawią się faceci 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

Podejście nr 2: Kolejne z niektórymi 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))

Alternatywnie ostatni krok cumsummożna zastąpić repeatfunkcjonalnością -

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

Podejście nr 3: Kolejne z przeważnie tylko 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

Możesz pracować z drugiej strony. Trzymaj licznik liczby niezerowych cyfr i przypisz je do elementu w tablicy. Jeśli widzisz 0, zresetuj licznik na 0

Edycja: jeśli po prawej stronie nie ma zera, musisz ponownie sprawdzić

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)

Możesz użyć różnicy między indeksami każdej pozycji i skumulowanych maks. Pozycji zerowych, aby określić odległość do poprzedniego zera. Można to zrobić do przodu i do tyłu. Minimalna odległość do przodu i do tyłu do poprzedniego (lub następnego) zera będzie najbliższa:

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)

wyniki:

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]

Specjalny przypadek, w którym na zewnętrznych krawędziach nie ma zer:

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]

działa również bez zer

[EDYCJA]  rozwiązania niepoliczalne ...

jeśli szukasz rozwiązania O (N), które nie wymaga numpy, możesz zastosować tę strategię za pomocą funkcji akumulacji z 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)

wynik:

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]

Jeśli nie chcesz korzystać z żadnej biblioteki, możesz ręcznie gromadzić odległości w pętli:

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)

wynik:

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]

Moją pierwszą intuicją byłoby użycie krojenia. Jeśli x może być normalną listą zamiast tablicy numpy, możesz użyć

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

jeśli numpy jest konieczne, możesz użyć

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

ale jest to mniej wydajne, ponieważ znajdujesz wszystkie zerowe lokalizacje po prawej stronie wartości, a następnie wyciągasz tylko pierwszą. Zdecydowanie lepszy sposób na zrobienie tego w numpy.

Edycja: Przepraszam, źle zrozumiałem. To da ci odległość do najbliższych zer - może to być po lewej lub po prawej stronie. Ale możesz użyć d_rightjako wyniku pośredniego. Nie obejmuje to jednak przypadku, w którym nie ma żadnego zera po prawej stronie.

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