python - 从 nd 数组中的 bin 索引查找封闭网格点的索引

Question

我有一个网格点的 nd 数组 G，然后是一个覆盖的 bin nd 数组 B，如下图所示。 G 的索引显示为蓝色，B 的索引显示为红色。给定 bin 的索引(红色)，我想找到封闭网格点的索引(蓝色)。

例如，给定 bin 索引 7，封闭的网格点索引将为 (1, 2, 5, 6)。

我正在寻找一种适用于 n 维数组而不仅仅是 2D 数组的解决方案。我感觉这应该是一个经常出现的问题，但在 numpy 中没有找到任何解决方案，并且正在努力为 n 维提出一个优雅的解决方案。

请注意，沿每个数组维度，B 的索引数量比 G 大 1。

<小时/>

一些测试代码

import numpy as np
G = np.arange(3*4).reshape(3, 4)
B = np.arange(4*5).reshape(4, 5)
idx = 7

Answer 1

您可以使用以下代码构建一个字典来关联索引。

from collections import defaultdict

#Map indices
d = defaultdict(list)
for n,row in enumerate(G):
    for i,idx in enumerate(row):
        for j in range(2): #check current and next row
            d[idx].append(B[n+j][i]) #current row
            try:  #next row (or next number) may not exist
                #next row
                #[i+1] is the index of the next number in the row
                d[idx].append(B[n+j][i+1]) 
            except IndexError:
                pass

这将创建

>>> d
defaultdict(list,
            {0: [0, 1, 5, 6],
             1: [1, 2, 6, 7],
             2: [2, 3, 7, 8],
             3: [3, 4, 8, 9],
             4: [5, 6, 10, 11],
             5: [6, 7, 11, 12],
             6: [7, 8, 12, 13],
             7: [8, 9, 13, 14],
             8: [10, 11, 15, 16],
             9: [11, 12, 16, 17],
             10: [12, 13, 17, 18],
             11: [13, 14, 18, 19]})

然后只需在字典的值中查找idx即可。

#Search value
idx = 7 #red bin
r = []
for k,v in d.items():
    if idx in v:
        r.append(k)

结果:

>>> print(r) #blue bins sorrounding idx
[1, 2, 5, 6]