python - 对角蛇填充数组

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

我做错了什么？ 附言您的回答可以是任何语言。

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 索引。

更准确地说，return 语句中的前三行计算对角线偏移量。第一行计算左上三角形中的对角线，第二行计算全长对角线以及右下三角形中的对角线；它也将这些计算为全长 - 第三行对此进行了更正。

最后两行在对角线内计数。两个中的第一个在右上方向，第二个在左下方向。请注意，对于从左边缘开始的所有对角线，右上角的偏移量等于 x 坐标。校正项 (np.clip ...) 适用于从底部边缘开始的对角线。类似地，如果我们从顶部边缘开始，左下角偏移量为 y，如果我们从右边缘开始，则需要更正。