python - 使用 NumPy 查找所有 n 维直线和对角线

Question

我想使用 NumPy 生成一个包含长度为 k 的 n 维数组的所有直线和对角线的列表。

---

以下面的长度为三的三维数组为例。

array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])

对于这种情况，我想获得以下所有类型的序列。对于任何给定的情况，我想获得每种类型的所有可能序列。对于每种情况，下面的括号中给出了所需序列的示例。

• 一维线
  • x 轴(0、1、2)
  • y 轴(0、3、6)
  • z 轴(0、9、18)
• 二维对角线
  • x/y 轴(0、4、8，2、4、6)
  • x/z 轴(0, 10, 20, 2, 10, 18)
  • y/z 轴(0、12、24，6、12、18)
• 3D 对角线
  • x/y/z 轴(0、13、26，2、13、24)

---

该解决方案应该被推广，以便它会生成一个数组的所有直线和对角线，而不管数组的维数或长度(它在所有维度上都是常数)。

Answer 1

此解决方案在 n

上推广

让我们将此问题改写为“查找索引列表”。

我们正在寻找形式的所有二维索引数组

array[i[0], i[1], i[2], ..., i[n-1]]

令n = arr.ndim

其中 i 是形状为 (n, k) 的数组

i[j] 中的每一个都可以是以下之一:

• 相同的索引重复 n 次，ri[j] = [j, ..., j]
• 正向序列，fi = [0, 1, ..., k-1]
• 后向序列，bi = [k-1, ..., 1, 0]

要求每个序列的形式为^(ri)*(fi)(fi|bi|ri)*$(用regex总结)。这是因为:

• 必须至少有一个fi所以“线”不是重复选择的点
• bi 不会出现在 fi 之前，以避免反转行

---

def product_slices(n):
    for i in range(n):
        yield (
            np.index_exp[np.newaxis] * i +
            np.index_exp[:] +
            np.index_exp[np.newaxis] * (n - i - 1)
        )

def get_lines(n, k):
    """
    Returns:
        index (tuple):   an object suitable for advanced indexing to get all possible lines
        mask (ndarray):  a boolean mask to apply to the result of the above
    """
    fi = np.arange(k)
    bi = fi[::-1]
    ri = fi[:,None].repeat(k, axis=1)

    all_i = np.concatenate((fi[None], bi[None], ri), axis=0)

    # inedx which look up every possible line, some of which are not valid
    index = tuple(all_i[s] for s in product_slices(n))

    # We incrementally allow lines that start with some number of `ri`s, and an `fi`
    #  [0]  here means we chose fi for that index
    #  [2:] here means we chose an ri for that index
    mask = np.zeros((all_i.shape[0],)*n, dtype=np.bool)
    sl = np.index_exp[0]
    for i in range(n):
        mask[sl] = True
        sl = np.index_exp[2:] + sl

    return index, mask

---

应用于您的示例:

# construct your example array
n = 3
k = 3
data = np.arange(k**n).reshape((k,)*n)

# apply my index_creating function
index, mask = get_lines(n, k)

# apply the index to your array
lines = data[index][mask]
print(lines)

array([[ 0, 13, 26],
       [ 2, 13, 24],
       [ 0, 12, 24],
       [ 1, 13, 25],
       [ 2, 14, 26],
       [ 6, 13, 20],
       [ 8, 13, 18],
       [ 6, 12, 18],
       [ 7, 13, 19],
       [ 8, 14, 20],
       [ 0, 10, 20],
       [ 2, 10, 18],
       [ 0,  9, 18],
       [ 1, 10, 19],
       [ 2, 11, 20],
       [ 3, 13, 23],
       [ 5, 13, 21],
       [ 3, 12, 21],
       [ 4, 13, 22],
       [ 5, 14, 23],
       [ 6, 16, 26],
       [ 8, 16, 24],
       [ 6, 15, 24],
       [ 7, 16, 25],
       [ 8, 17, 26],
       [ 0,  4,  8],
       [ 2,  4,  6],
       [ 0,  3,  6],
       [ 1,  4,  7],
       [ 2,  5,  8],
       [ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 13, 17],
       [11, 13, 15],
       [ 9, 12, 15],
       [10, 13, 16],
       [11, 14, 17],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 22, 26],
       [20, 22, 24],
       [18, 21, 24],
       [19, 22, 25],
       [20, 23, 26],
       [18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]])

---

另一组很好的测试数据是np.moveaxis(np.indices((k,)*n), 0, -1)，它给出了一个数组，其中每个值都是它自己的索引

---

我之前已经解决了这个问题来实现 higher dimensional tic-tac-toe