Найти расстояние до ближайшего нуля в массиве 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 к спасению для линейного времени в векторизации (до того, как придут парни из 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

Вы могли бы работать с другой стороны. Держите счетчик на количество переданных ненулевых цифр и присвойте его элементу в массиве. Если вы видите 0, сбросьте счетчик на 0

Редактировать: если справа нет нуля, то вам нужна еще одна проверка

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)

Вы можете использовать разницу между индексами каждой позиции и совокупным максимумом нулевых позиций, чтобы определить расстояние до предыдущего нуля. Это может быть сделано вперед и назад. Минимум между прямым и обратным расстоянием до предыдущего (или следующего) нуля будет ближайшим:

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)

Результаты:

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]

Особый случай, когда на внешних ребрах нет нулей:

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]

также работает без нулей вообще

[РЕДАКТИРОВАТЬ]  non-numpy решения ...

если вы ищете решение O (N), которое не требует numpy, вы можете применить эту стратегию, используя функцию накопления из 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)

вывод:

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]

Если вы не хотите использовать какую-либо библиотеку, вы можете накапливать расстояния вручную в цикле:

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)

вывод:

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]

Моей первой интуицией было бы использовать нарезку. Если x может быть обычным списком вместо массива, вы можете использовать

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

если нужна numpy, вы можете использовать

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

но это менее эффективно, потому что вы находите все нулевые позиции справа от значения, а затем вытаскиваете только первое. Почти наверняка лучший способ сделать это в NumPy.

Изменить: я извиняюсь, я неправильно понял. Это даст вам расстояние до ближайших нулей - может ли оно быть слева или справа. Но вы можете использовать d_rightкак промежуточный результат. Это не охватывает крайний случай отсутствия нуля вправо.

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