生成3D阵列中坐标的邻居索引列表

时间:2017-09-09 20:50:27

标签: python itertools

我想生成3D数组中给定点的某些邻居的索引。更简单的方法是手动编写它们。这里有一个示例,其中一个方向将具有一系列可能的值,但另外两个可能只有0或1.我在这个示例中使用的范围从-3到3而不是0使它太长,但我用-6到6。

def gen_neigh_indices(i, j, k):
    for o, p, q in ((0, 0, 1), (0, 0, 2), (0, 0, 3),
                    (0, 1, 1), (0, 1, 2), (0, 1, 3),
                    (1, 0, 1), (1, 0, 2), (1, 0, 3),
                    (1, 1, 1), (1, 1, 2), (1, 1, 3),
                    (0, 0, -1), (0, 1, -2), (0, 1, -3),
                    ...
                    (0, 1, 0), (0, 2, 0), (0, 3, 0),
                    ...
                    (0, -1, 0), (0, -2, 0), (0, -3, 0),
                    ...
                    (1, 0, 0), (2, 0, 0), (3, 0, 0),
                    ...
                    (-1, 1, 1), (-2, 1, 1), (-3, 1, 1)):
        yield (i+o, j+p, k+q)

我想过使用itertools函数但到目前为止我还不知道如何更改返回值的顺序而不是手动,所以我写了3个收益。

from itertools import product

def get_neigh_indices(i, j, k):
    for o in range(-3, 4):
        if not o: continue
        for p, q in product(range(2), repeat=2):
            yield (i+p, j+q, k+o)
            yield (i+p, j+o, k+q)
            yield (i+o, j+p, k+q)

有没有办法只写一个yield?我不关心所生成元素的顺序,只是它提供了所有可能的组合。

0 个答案:

没有答案