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

Question

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

---

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

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]]])

对于此案例，我想获取以下所有类型的序列。 对于任何给定的案例，我都希望获取每种类型的所有可能序列。 下面为每种情况给出所需序列的示例。

- 1D线段

- x轴 (0, 1, 2) - y轴 (0, 3, 6) - z轴 (0, 9, 18)

- 2D对角线

- 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)

解决方案应该是一般化的，因此它将为一个数组生成所有行和对角线，而不考虑数组的维数或长度（在所有维度中均为常数）。

- 2Cubed

---

2个突出的块对于2D对角线有什么重要性？中央对角线将有3个项目；您将反对角线包含在同一个元组中，但以某种方式分开。您是否也对偏移对角线感兴趣？ - hpaulj

你想生成(0, 1, 2)和(2, 1, 0)吗？ - Eric

你是否正在尝试实现 这个？ - Eric

@hpaulj 他们旨在展示我对两个方向的对角线都很感兴趣。 - 2Cubed

@Eric 我不想生成 (0, 1, 2) 和 (2, 1, 0) 这两个 - 只需要其中一个即可。从概念上讲，我想实现的是那个游戏的逻辑，但是以一般化的格式（适用于任何大小或维数）。 - 2Cubed

那个游戏泛化了 k 但不是 n。我在这里的答案两者都进行了泛化。 - Eric

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)*$（使用正则表达式进行总结）。这是因为：

• 必须至少有一个 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)，它提供了一个数组，其中每个值都是其自身的索引。

---

我以前曾经解决过这个问题，以实现高维井字棋游戏。

- Eric

Answer 2

In [1]: x=np.arange(27).reshape(3,3,3)

选择单个行很容易：

In [2]: x[0,0,:]
Out[2]: array([0, 1, 2])
In [3]: x[0,:,0]
Out[3]: array([0, 3, 6])
In [4]: x[:,0,0]
Out[4]: array([ 0,  9, 18])

您可以使用索引列表迭代维度：

In [10]: idx=[slice(None),0,0]
In [11]: x[idx]
Out[11]: array([ 0,  9, 18])
In [12]: idx[2]+=1
In [13]: x[idx]
Out[13]: array([ 1, 10, 19])

查看np.apply_along_axis的代码，了解它如何实现此类迭代。

重塑和分割也可以生成行列表。对于某些维度，这可能需要进行转置：

In [20]: np.split(x.reshape(x.shape[0],-1),9,axis=1)
Out[20]: 
[array([[ 0],
        [ 9],
        [18]]), array([[ 1],
        [10],
        [19]]), array([[ 2],
        [11],
        ...

np.diag可以从2D子数组中获取对角线元素。

In [21]: np.diag(x[0,:,:])
Out[21]: array([0, 4, 8])
In [22]: np.diag(x[1,:,:])
Out[22]: array([ 9, 13, 17])
In [23]: np.diag?
In [24]: np.diag(x[1,:,:],1)
Out[24]: array([10, 14])
In [25]: np.diag(x[1,:,:],-1)
Out[25]: array([12, 16])

探索 np.diagonal ，可直接应用于三维数组。也可以使用 range 和 arange 直接索引数组，例如： x[0,range(3),range(3)]。

据我所知，没有一个函数可以遍历所有这些替代方案。由于返回的数组维数可能不同，在已编译的numpy代码中制作这样的功能几乎没有意义。因此，即使有一个函数，它也会按照我概述的方式遍历这些替代方案。

==============

所有的一维线路

x.reshape(-1,3)
x.transpose(0,2,1).reshape(-1,3)
x.transpose(1,2,0).reshape(-1,3)

y/z对角线和反对角线

In [154]: i=np.arange(3)
In [155]: j=np.arange(2,-1,-1)
In [156]: np.concatenate((x[:,i,i],x[:,i,j]),axis=1)
Out[156]: 
array([[ 0,  4,  8,  2,  4,  6],
       [ 9, 13, 17, 11, 13, 15],
       [18, 22, 26, 20, 22, 24]])

- hpaulj

---

这些函数并未获取每个情况下的所有序列。例如，x[0,0,:]仅检索第一行，而它应该获取[0, 1, 2]、[3, 4, 5]、……、[21, 22, 23]、[24, 25, 26]。我问题中括号内的序列只是所需输出的示例。 - 2Cubed

我谈到了迭代需要获取所有切片的各种索引。 - hpaulj

Answer 3

np.einsum可以用来构建所有这些类型的表达式，例如：

# 3d diagonals
print(np.einsum('iii->i', a))
# 2d diagonals
print(np.einsum('iij->ij', a))
print(np.einsum('iji->ij', a))

- Eelco Hoogendoorn

---

但这是求和，不仅仅是切片。 - Eric

它仅对出现在左侧而不在右侧的索引求和。 - Eelco Hoogendoorn