Find The Nearest Nonzero Element And Corresponding Index In A 2d NumPy Array

Let's say I have a 2d array named array and a 2d index:(x,y) and I want to get the nearest nonzero element to (x,y), and get the corresponding index of the element in np.nonzero(ar

Solution 1:

Approach #1 : Here's one approach -

def nearest_nonzero_idx(a,x,y):
    idx = np.argwhere(a)

    # If (x,y) itself is also non-zero, we want to avoid those, so delete that
    # But, if we are sure that (x,y) won't be non-zero, skip the next step
    idx = idx[~(idx == [x,y]).all(1)]

    return idx[((idx - [x,y])**2).sum(1).argmin()]

Sample run -

In [64]: a
Out[64]: 
array([[0, 0, 1, 1, 0, 1, 1],
       [0, 0, 1, 0, 0, 0, 1],
       [1, 1, 0, 0, 0, 0, 0],
       [0, 1, 1, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0],
       [1, 1, 0, 0, 1, 1, 0],
       [0, 1, 0, 0, 1, 0, 1],
       [1, 0, 0, 1, 1, 1, 0]])

In [65]: x,y =(3,5)

In [66]: nearest_nonzero(a,x,y)
Out[66]: array([5, 5])

Approach #2 : Here's another approach with focus on performance that avoids the second step from previous approach of skipping the (x,y) point from the array of non-zero indices by temporarily setting the point to 0, getting those non-zero indices and then setting back the original value into it. Also, we could use np.nonzero to store the row, col indices and then perform the distance computing with those.

Thus, the implementation would be -

def nearest_nonzero_idx_v2(a,x,y):
    tmp = a[x,y]
    a[x,y] = 0
    r,c = np.nonzero(a)
    a[x,y] = tmp
    min_idx = ((r - x)**2 + (c - y)**2).argmin()
    return r[min_idx], c[min_idx]

Runtime test

In [110]: a
Out[110]: 
array([[3, 2, 3, 3, 0, 2, 4, 2, 1],
       [0, 3, 4, 3, 4, 3, 3, 2, 0],
       [1, 3, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 2, 0, 0, 2, 0, 0, 2],
       [3, 0, 0, 0, 0, 0, 0, 0, 1],
       [0, 0, 2, 2, 4, 4, 3, 4, 3],
       [2, 2, 2, 1, 0, 0, 1, 1, 1],
       [3, 4, 3, 1, 0, 4, 0, 4, 2]])

In [111]: x,y =(3,5)

In [112]: nearest_nonzero_idx(a,x,y)
Out[112]: array([1, 5])

In [113]: nearest_nonzero_idx_v2(a,x,y)
Out[113]: (1, 5)

In [114]: %timeit nearest_nonzero_idx(a,x,y)
10000 loops, best of 3: 23.1 µs per loop

In [115]: %timeit nearest_nonzero_idx_v2(a,x,y)
100000 loops, best of 3: 4.85 µs per loop