Diagonal snake filling array

Question

Python 3.7. I'm trying to fill multidimensional array (n*m size) in diagonal-snake pattern:

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

I have a function for n x n size and it works fine for it. But for n x m size it returns:

1 3  4  10 14

2 5  9  15 20

6 8  16 19 19

7 17 18 20 21

My code:

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

What am I doing wrong? P.S. Your answer can be in any language.

Answer 1

Here is a vectorized solution:

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

Demo:

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

Explainer:

The function tr computes the number of elements in a triangle which is more or less half a square (a tiny bit more because we include the diagonal). This is used in snake to compute the offset of each diagonal; diagonals are indexed by x+y.

More precisely, the first three lines in the return statement compute the diagonal offset. The first line counts diagonals in the top left triangle, the second line counts full length diagonals and also those in the bottom right triangle; it counts those also as full length - the third line corrects for that.

The last two lines count within diagonals. The first of the two in top-right direction, the second in bottom-left direction. Observe, that the top-right offset is equal to the x coordinate for all diagonals that start at the left edge. The correction term (np.clip ...) is for diagonals starting at the bottom edge. Similarly bottom-left offsets are y if we start at the top edge and require a correction if we start at the right edge.