在 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 loop 按照下面提到的代码进行此练习。

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循环快速解决这个问题的方法？

Answer 1

我建议使用类似这种方法，它在矩阵元素的数量上是线性的（考虑到您支付的费用不能少于至少扫描一次矩阵的所有元素）。本质上，我将相邻数字的列表分组为一组，然后在此基础上计算邻居条目。

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)

我在一个包含 100 万个不同数字 (np.random.randint(0,10,(1000,1000))) 的随机矩阵上对其进行了测试，耗时 1.61 秒。

---

更新：

同样的方法可以用于 3D 数组。代码如下：

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)

我还在具有 10 个不同数字 (np.random.randint(0,10,(100,100,100))) 的 1M 元素的随机矩阵上测试了此函数，耗时 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

由于这个功能是通用的，它需要做更多的工作，实际上它比以前的慢。在第一次测试的相同 2D 矩阵上需要 6.71 秒，而在第二次测试的 3D 矩阵上需要 7.96 秒。

---

更新 2：

我用更快（我希望也更简单）的版本更新了 2D 和 3D 矩阵的代码。如果没有关于矩阵内数字位置的其他限制，则无法在不扫描所有矩阵单元的情况下发现所有颜色：使用 for 循环，我们目前正在这样做。您可以用来完成任务的每个功能都将在内部扫描整个矩阵（至少）。顺便说一句，我并不是说这是最快的解决方案，因为另一种解决方案可以是使用 cython 或本机 numpy 代码（如果存在）。

Answer 2

您可以使用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])}

注意我修剪了边缘；如果您想包括环绕邻居或做一些更细微的事情，您可以详细说明修剪线。