在2D或3D数组中查找相邻数字的快速方法

Question

我有以下的二维数组

regions = array([[3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4],
                 [3, 3, 3, 3, 8, 8, 8, 8, 8, 4, 4, 4, 4],
                 [3, 3, 3, 3, 8, 8, 8, 8, 8, 4, 4, 4, 4],
                 [3, 3, 3, 3, 8, 8, 8, 8, 8, 4, 4, 4, 4],
                 [3, 6, 6, 6, 8, 8, 8, 8, 8, 7, 7, 7, 4],
                 [3, 6, 6, 6, 8, 8, 8, 8, 8, 7, 7, 7, 4],
                 [3, 6, 6, 6, 6, 8, 8, 8, 7, 7, 7, 7, 4],
                 [3, 6, 6, 6, 6, 2, 2, 2, 7, 7, 7, 7, 4],
                 [5, 6, 6, 6, 6, 2, 2, 2, 7, 7, 7, 7, 1],
                 [5, 6, 6, 6, 6, 2, 2, 2, 7, 7, 7, 7, 1],
                 [5, 6, 6, 6, 6, 2, 2, 2, 7, 7, 7, 7, 1],
                 [5, 5, 5, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1],
                 [5, 5, 5, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1]])

我希望能够为所有的数字找到相邻的数字。例如， 3 是 4,5,6,8 的相邻数字。目前我正在使用以下代码来完成这个练习。

for循环

numbers = scipy.unique(regions)
for i in numbers:
    index = i-1
    slices = scipy.ndimage.find_objects(regions)
    sub_im = regions[slices[index]]
    im = sub_im == i
    neighbors = scipy.ndimage.binary_dilation(input=im, structure=disk(1))
    neighbors = neighbors*sub_im
    neighbors_list = scipy.unique(neighbors)[1:] - 1
    print (neighbors_list)

问题是我不想使用for循环，因为我的区域数组有数百万个。有没有快速解决此问题的方法，而不需要使用for循环？

- ZAK

---

你可以迭代2x2的窗口；如果窗口中有3，则保存窗口中不是3的所有数字。搜索numpy滑动/移动窗口以获得高效的窗口迭代器/生成器/步幅技巧。这值得一试以进行比较。 - wwii

感谢您的回复@wwii。您有没有任何简单的链接可以解释这种方法？ - ZAK

Answer 1

您可以使用numpy，使用np.roll()，np.where()和np.unique()来完成此操作。其想法是将每个条目附加到其四个邻居，并从这些五个成员的列表中提取唯一的成员（条目及其四个邻居）。以下是应该能够澄清的实现：

# make a 3d array with the matrix entry and its four neighbors
neighbor_array = np.array([regions,
                          np.roll(regions,+1,axis=0),
                          np.roll(regions,-1,axis=0),
                          np.roll(regions,+1,axis=1),
                          np.roll(regions,-1,axis=1),
                          ])
# if you want neighbors to include wraparounds, use above; if not, prune
neighbor_array_pruned = neighbor_array[:,1:-1,1:-1]
# reshape to a 2d array entries x neighbors
neighbor_list = np.reshape(neighbor_array_pruned,[5,-1]).T
# get uniques into a dictionary 
neighbor_dict = {}
for num in np.unique(regions):
    neighbor_dict[num] = np.unique(neighbor_list[np.where(neighbor_list[:,0]==num)])

这将生成neighbor_dict:

{1: array([1, 2, 7]),
 2: array([1, 2, 5, 6, 7, 8]),
 3: array([3, 6, 8]),
 4: array([4, 7, 8]),
 5: array([2, 5, 6]),
 6: array([2, 3, 5, 6, 8]),
 7: array([1, 2, 4, 7, 8]),
 8: array([2, 3, 4, 6, 7, 8])}

请注意，我修剪了边缘; 如果您想包括环绕邻居或执行更复杂的操作，则可以详细说明修剪行。

- muskrat

---

谢谢您的回复。代码运行良好，但是最后两行中的for循环正在减慢代码速度。它可以被任何数组操作所替代吗？ - ZAK

我尝试过去掉“for循环”，但目前还没有成功。 - ZAK

你的最终结果neighbor_dict显示了不完整的值。例如，第一个条目是1：array（[1,2,7]），而应该是1：array（[1,2,7,4]），因为regions数组中4也与1相连。有什么遗漏吗？@muskrat - ZAK

根据代码注释所讨论的，修剪步骤会忽略边缘；如果要包含环绕，则将其注释掉。如果不想要环绕，则可以使用NaN或0等填充边缘，然后在处理后忽略它们。请点赞或接受答案。 - muskrat

让我们在聊天中继续这个讨论。 - ZAK

Answer 2

我建议使用这种方法，它与矩阵元素数量呈线性关系（考虑到您不能只扫描矩阵的部分元素）。基本上，我将相邻数字列表分组成集合，然后在此基础上计算相邻条目。

import numpy as np
import collections

def adj_to_neighbor_dict(adj):
    assert hasattr(adj, "__iter__")

    neighbor_dict = collections.defaultdict(lambda: set())
    for i,j in adj:
        if i == j:
            continue
        neighbor_dict[i].add(j)
        neighbor_dict[j].add(i)
    return neighbor_dict

def get_neighbors_2d(npmatrix):
    assert len(npmatrix.shape) == 2
    I, J = range(npmatrix.shape[0]-1), range(npmatrix.shape[1]-1)
    adj_set = set(
        (npmatrix[i,j], npmatrix[i+1,j])
        for i in I
        for j in J
    ) | set(
        (npmatrix[i,j], npmatrix[i,j+1])
        for i in I
        for j in J
    )
    return adj_to_neighbor_dict(adj_set)

我对一个包含1百万个元素和10个不同数字的随机矩阵进行了测试（np.random.randint(0,10,(1000,1000))），用时1.61秒。

---

更新：

同样的方法也可以用于三维数组。实现代码如下：

def get_neighbors_3d(npmatrix):
    assert len(npmatrix.shape) == 3
    I, J, K = range(npmatrix.shape[0]-1), range(npmatrix.shape[1]-1), range(npmatrix.shape[2]-1)
    adj_set = set(
        (npmatrix[i,j,k], npmatrix[i+1,j,k])
        for i in I
        for j in J
        for k in K
    ) | set(
        (npmatrix[i,j,k], npmatrix[i,j+1,k])
        for i in I
        for j in J
        for k in K
    ) | set(
        (npmatrix[i,j,k], npmatrix[i,j,k+1])
        for i in I
        for j in J
        for k in K
    )
    return adj_to_neighbor_dict(adj_set)

我还在一个由1百万个元素和10个不同数字组成的随机矩阵上测试了这个函数（np.random.randint(0,10,(100,100,100))），运行时间为2.60秒。

---

我还会提出一个通用的解决方案，不基于np.array的形状。

def npmatrix_shape_iter(shape):
    num_dimensions = len(shape)
    last_dimension = num_dimensions-1
    coord = [0] * num_dimensions
    while True:
        yield tuple(coord)
        coord[last_dimension] += 1
        for i in xrange(last_dimension, 0, -1):
            if coord[i] < shape[i]:
                break
            coord[i] = 0
            coord[i-1] += 1
        # end condition: all the dimensions have been explored
        if coord[0] >= shape[0]:
            break

def adj_position_iter(tpl):
    new_tpl = list(tpl)
    for i in xrange(len(tpl)):
        new_tpl[i] += 1
        yield tuple(new_tpl)
        new_tpl[i] -= 1

def get_neighbors(npmatrix):
    neighbors = set(
        (npmatrix[tpl], npmatrix[adj_tpl])
        for tpl in npmatrix_shape_iter(tuple(np.array(npmatrix.shape)-1))
        for adj_tpl in adj_position_iter(tpl)
    )
    neighbor_dict = collections.defaultdict(lambda: [])
    for i,j in neighbors:
        if i == j:
            continue
        neighbor_dict[i].append(j)
        neighbor_dict[j].append(i)
    return neighbor_dict

由于这个函数是通用的，所以需要更多的工作来完成，实际上它比之前的函数要慢。在第一个测试的同样的二维矩阵上，它需要6.71秒，而在第二个测试的三维矩阵上，它需要7.96秒。

---

更新2：

我用更快（也希望更容易）的版本更新了2D和3D矩阵的代码。如果没有其他关于数字在矩阵中位置的限制，那么没有办法在不扫描所有矩阵单元的情况下发现所有颜色：我们目前正在使用for循环来做这件事。你可以使用的每个函数都将在内部扫描整个矩阵（至少如此）。顺便说一句，我并不是说这是最快的解决方案，因为另一种替代方案可能是使用cython或本地numpy代码（如果存在的话）。

- Roberto Trani

---

感谢您的回复，@Roberto Trani。我能在三维数组中使用这种方法吗？它对二维数组运行良好。 - ZAK

谢谢，@Roberto Trani。我已经尝试了您更新的3D数组代码，它运行良好。我在想我们是否有可能替换数组中的tpl和adj_tpl for循环，因为当我尝试(300,300,100)矩阵时，这两个for循环会大大增加代码时间。平均时间为25秒。 - ZAK