Python 在2d NumPy数组中查找最近的非零元素和相应索引

假设我有一个名为

array

的2d数组和一个2d索引：

（x，y）

，我想得到距离

（x，y）

最近的非零元素，并在

np.nonzero（array）

方法中得到元素的相应索引-

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

样本运行-

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

方法#2:这里是另一种注重性能的方法，它避免了前一种方法的第二步，即从非零索引数组中跳过

（x，y）

点，方法是临时将该点设置为

0

，获取这些非零索引，然后将原始值设置回该点。此外，我们还可以使用

np.nonzero

存储

行、列

索引，然后使用这些索引执行距离计算

因此，执行工作将是非常重要的-

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]

运行时测试

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

方法#1:这里有一种方法-

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

样本运行-

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

方法#2:这里是另一种注重性能的方法，它避免了前一种方法的第二步，即从非零索引数组中跳过

（x，y）

点，方法是临时将该点设置为

0

，获取这些非零索引，然后将原始值设置回该点。此外，我们还可以使用

np.nonzero

存储

行、列

索引，然后使用这些索引执行距离计算

因此，执行工作将是非常重要的-

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]

运行时测试

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