Знайдіть відстань до найближчого нуля в масиві 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]

також працює зовсім без нулів

[EDIT]  нечислові рішення ...

якщо ви шукаєте рішення 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)]

якщо нумінг необхідний, то ви можете використовувати

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

але це менш ефективно, оскільки ви знаходите всі нульові місця праворуч від значення, а потім витягуєте лише перше. Майже напевно кращий спосіб зробити це в нуме.

Редагувати: Вибачте, я зрозумів неправильно. Це дасть вам відстань до найближчих нулів - це може бути зліва або справа. Але ви можете використовувати 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])