由 2D/3D NumPy 数组表示的边界距离

Question

我有大型 2D/3D NumPy 二进制值数组。 1 代表边界，0 代表区域。附带一个步长数组，表示每个维度的步长大小。

我正在寻找一个高效的程序，它可以找到（其中一个）与给定元素最近的边界元素。距离是欧式距离。

2D 示例：

import seaborn as sns
import numpy as np

step_size = np.array([5,5])  # size of step in each dimension
arr = np.array([[0,0,0,0,0,0,0,0,1,0],[1,0,0,0,0,0,0,1,0,0],[1,0,0,0,0,0,0,1,0,0],[0,1,0,0,0,0,1,0,0,0],[0,0,1,0,0,0,0,1,0,0],[0,0,1,1,1,1,1,0,0,0]])
sns.heatmap(arr, annot=True, cbar=False, linewidths=.3)

a = (0,2)  # a given element index
b = (1,0)  # nearest boundary element index, which is to be found by the program

a_coor = np.multiply(np.array(a), step_size)  # array([0,10])
b_coor = np.multiply(np.array(b), step_size)  # array([5,0])

distance = np.linalg.norm(a_coor-b_coor)  # 11.18

Answer 1

绝对是最近邻搜索的工作：

import seaborn as sns
import numpy as np
from sklearn.neighbors import KDTree
from matplotlib import pyplot as plt

step_size = np.array([5,5])  # size of step in each dimension
arr = np.array([[0,0,0,0,0,0,0,0,1,0],[1,0,0,0,0,0,0,1,0,0],[1,0,0,0,0,0,0,1,0,0],[0,1,0,0,0,0,1,0,0,0],[0,0,1,0,0,0,0,1,0,0],[0,0,1,1,1,1,1,0,0,0]])
sns.heatmap(arr, annot=True, cbar=False, linewidths=.3)

# get boundary pts (+0.5 to center for plot)
boundary = np.column_stack(np.nonzero(arr)) + 0.5

# create tree
tree = KDTree(boundary)

# get zero points to test for
zeros = np.column_stack(np.nonzero(~arr)) + 0.5

# get nearest neighbour boundary point for each zero
distance, index = tree.query(zeros)

# plot the results
for i, pt in enumerate(zeros):
    plt.gca().plot([pt[1], boundary[index[i,0]][1]], [pt[0], boundary[index[i,0]][0]], 'r')

KDTree 可以轻松计算k 邻居，并返回原始树中结果的欧几里得distance 和index。一个非常有用的工具。见下图：

还有第二近邻的结果，将k=2 传递给query 并绘制：

# plot the results
colors = ['r', 'b']
for i, pt in enumerate(zeros):
    for k in range(index.shape[1]):
        plt.gca().plot([pt[1], boundary[index[i,k]][1]], [pt[0], boundary[index[i,k]][0]], colors[k])

Answer 2

您可以找到所有1s 的位置，获取到给定位置的欧几里得距离，并在结果上使用argpartition 返回最小值：

def nearest_boundary(x, loc):
    ones = np.c_[np.where(arr == 1)]
    dist = ((ones - loc)**2).sum(1)
    return ones[dist.argpartition(0)[0]]

---

一些例子：

nearest_boundary(arr, (0,2))
# array([1, 0], dtype=int64)

nearest_boundary(arr, (2,4))
# array([3, 6], dtype=int64)

使用 3D 数组：

np.random.seed(2)
arr = np.random.choice([0,1], p=[0.8,0.2], size=(3,5,4))

array([[[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [1, 0, 1, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[1, 0, 1, 0],
        [0, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 1]]])

nearest_boundary(arr, (0,3,0))
# array([0, 4, 0], dtype=int64)