对角线蛇形填充数组

Question

我正在尝试使用Python 3.7填充多维数组（n*m大小）以对角蛇形模式：

1   3   4   10  11  21
2   5   9   12  20  22
6   8   13  19  23  30
7   14  18  24  29  31
15  17  25  28  32  35
16  26  27  33  34  36

我有一个针对 n x n 大小的函数，它可以正常工作。但是对于 n x m 大小，它会返回：

1 3  4  10 14

2 5  9  15 20

6 8  16 19 19

7 17 18 20 21

我的代码：

def method1(i, j, n, m):
    num = i+j
    summ = num * (num + 1) >> 1
    s = n * m
    if num > n-1:
        t = 2*(n-1) - (i+j) + 1
        s -= t*(t+1) >> 1

    if num & 1:
        if num > n-1:
            return s + (n-i)
        else:
            return summ + j+1
    if num > n-1:
        return s + (n-j)
    else:
        return summ + i+1

for i in range(n):
    for j in range(m):
        print(method1(i, j, n, m), end=" ")
    print('\n')

我做错什么了吗？ P.S.你的回答可以用任何语言。

- Neuron

Answer 1

这里是向量化的解决方案：

def tr(z):
    return z*(z+1)//2

def snake(Y, X):
    y, x = np.ogrid[:Y, :X]
    mn, mx = np.minimum(X, Y), np.maximum(X, Y)
    return (1 + tr(np.clip(x+y, None, mn))
            + mn * np.clip(x+y - mn, 0, None)
            - tr(np.clip(x+y - mx, 0, None))
            + ((x+y) & 1) * (x - np.clip(x+y + 1 - Y, 0, None))
            + ((x+y + 1) & 1) * (y - np.clip(x+y + 1 - X, 0, None)))

示例:

>>> snake(7, 3)
array([[ 1,  3,  4],
       [ 2,  5,  9],
       [ 6,  8, 10],
       [ 7, 11, 15],
       [12, 14, 16],
       [13, 17, 20],
       [18, 19, 21]])
>>> snake(2, 4)
array([[1, 3, 4, 7],
       [2, 5, 6, 8]])

解释：

tr 函数计算三角形中的元素数量，该三角形几乎是一个正方形的一半（因为我们包括对角线，所以多一点）。这在 snake 中用于计算每个对角线的偏移量；对角线由 x+y 进行索引。

更精确地说，返回语句中的前三行计算对角线偏移量。第一行计算左上三角形中的对角线数，第二行计算完整对角线数以及底部右侧三角形中的对角线数量；它也将那些对角线视为完整长度 - 第三行进行了修正。

最后两行计算对角线内部的数量。前面一行是朝右上方的方向，后面一行是朝左下方的方向。请注意，对于从左边开始的所有对角线，顶部右侧的偏移量等于 x 坐标。校正项 (np.clip ...) 用于从底边开始的对角线。同样，如果我们从右边开始，则底部左侧偏移为 y 并需要进行修正。

- Paul Panzer

---

这很令人印象深刻，但我不知道你是如何让它工作的！虽然不确定能否简洁地解释清楚。 - jdehesa

对我来说很难，但探索起来很有用。谢谢！ - Neuron

@jdehesa，EAMax我添加了一些解释。希望有所帮助。 - Paul Panzer

非常好的答案，而且解释得非常清楚！在理解不同部分方面，打印返回函数的各个加数对我很有帮助。 - pietroppeter

我认为这个答案比我的好多了（尤其是现在有了解释）。@EAMax，随意更改接受的答案！ - pietroppeter

Answer 2

编辑：

这里有一个基本相同的算法版本，但没有任何循环：

def snake_matrix(n):
    # Make sequences: [0, 0, 1, 0, 1, 2, 0, 1, 2, 3, ...]
    i = np.arange(n)
    c = np.cumsum(i)
    reps = np.repeat(c, i + 1)
    seqs = np.arange(len(reps)) - reps
    # Make inverted sequences: [0, 1, 0, 2, 1, 0, 3, 2, 1, 0, ...]
    i_rep = np.repeat(i, i + 1)
    seqs_inv = i_rep - seqs
    # Select sequences for row and column indices
    seq_even_mask = (i_rep % 2 == 0)
    # Row inverts even sequences
    row = np.where(seq_even_mask, seqs, seqs_inv)
    # Column inverts odd sequences
    col = np.where(~seq_even_mask, seqs, seqs_inv)
    # Mirror  for lower right corner
    row = np.concatenate([row, n - 1 - row[len(row) - n - 1::-1]])
    col = np.concatenate([col, n - 1 - col[len(col) - n - 1::-1]])
    m = np.empty((n, n), dtype=int)
    m[row, col] = np.arange(n * n)
    return m

有趣的是，在进行了几次基准测试后，似乎根据大小而定，这可能比之前的算法更快。

---

这里有另一种使用NumPy的解决方案。我不知道是否有其他方法可以使其更好（没有循环，或者在这种情况下列表推导式），但至少它不会遍历每个单独的元素。不过，这个解决方案只适用于方阵。

import numpy as np

def snake_matrix(n):
    # Sequences for indexing top left triangle: [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]...
    seqs = [np.arange(i + 1) for i in range(n)]
    # Row indices reverse odd sequences
    row = np.concatenate([seq if i % 2 == 0 else seq[::-1] for i, seq in enumerate(seqs)])
    # Column indices reverse even sequences
    col = np.concatenate([seq if i % 2 == 1 else seq[::-1] for i, seq in enumerate(seqs)])
    # Indices for bottom right triangle "mirror" top left triangle
    row = np.concatenate([row, n - 1 - row[len(row) - n - 1::-1]])
    col = np.concatenate([col, n - 1 - col[len(col) - n - 1::-1]])
    # Make matrix
    m = np.empty((n, n), dtype=int)
    m[row, col] = np.arange(n * n)
    return m

print(snake_matrix(6))

输出：

[[ 0  2  3  9 10 20]
 [ 1  4  8 11 19 21]
 [ 5  7 12 18 22 29]
 [ 6 13 17 23 28 30]
 [14 16 24 27 31 34]
 [15 25 26 32 33 35]]

关于这种枚举的更多信息可以在OEIS A319571序列中找到（尽管该序列是针对无限网格的一般序列，但在这种情况下，您将有一个从左上角开始的枚举和另一个从右下角开始的枚举）。

- jdehesa

---

感谢您的回复。非常有帮助。 - Neuron

@EAMax 很高兴能帮到你。我添加了一个没有循环的版本，只是为了好玩，尽管它并不总是比另一个更快。无论如何，正如我所说，这个算法仅适用于方阵，因此它仍然不是一个完整的解决方案。 - jdehesa

Answer 3

不清楚您做错了什么，但是以下代码应该可以正常工作：

import numpy as np

n = 4
m = 5

x, y = (0, 0)
ux, uy = (1, -1)

a = np.zeros((n, m))
for i in range(n*m):
  print((x, y), i+1)
  a[x, y] = i + 1
  x, y = (x + ux, y + uy)
  if y == m:
    print('right side')  # including corner
    y = m - 1
    x += 2
  elif x == n:
    print('bottom side')  # including corner
    x = n - 1
    y += 2
  elif x == -1:
    print('top side')
    x = 0
  elif y == -1:
    print('left side')
    y = 0
  else:
    continue
  ux, uy = -ux, -uy
print(a)

输出：

(0, 0) 1
left side
(1, 0) 2
(0, 1) 3
top side
(0, 2) 4
(1, 1) 5
(2, 0) 6
left side
(3, 0) 7
(2, 1) 8
(1, 2) 9
(0, 3) 10
top side
(0, 4) 11
(1, 3) 12
(2, 2) 13
(3, 1) 14
bottom side
(3, 2) 15
(2, 3) 16
(1, 4) 17
right side
(2, 4) 18
(3, 3) 19
bottom side
(3, 4) 20
right side
[[ 1.  3.  4. 10. 11.]
 [ 2.  5.  9. 12. 17.]
 [ 6.  8. 13. 16. 18.]
 [ 7. 14. 15. 19. 20.]]

为了写这篇文章，画图会很有帮助。

- pietroppeter