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.

Answer 2

EDIT:

Here is a version of basically the same algorithm but without any loops:

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

Interestingly, after a couple of benchmarks it seems that depending on the size this may or may not be faster than the previous algorithm.

---

Here is another solution with NumPy. I don't know if there is any other way to make this better (without loops, or in this case list comprehensions), but at least it does not loop over every single element. This one only works for square matrices though.

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

Output:

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

There is some more information about this kind of enumeration in OEIS A319571 sequence (although that refers to the general sequence for an infinite grid, in this case you would have one enumeration starting at the top left and another at the bottom right).

Answer 3

not clear what you are doing wrong, but the following code should work:

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)

output:

(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.]]

To write this, it helped a lot to draw a diagram.