Диагональный массив змей

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.

Точнее, первые три строки в операторе возврата вычисляют смещение по диагонали. Первая строка подсчитывает диагонали в верхнем левом треугольнике, вторая строка подсчитывает диагонали полной длины, а также в нижнем правом треугольнике; он также считает их полной длиной - третья строка исправляет это.

Последние две строки считаются в диагонали. Первый из двух в верхнем правом направлении, второй в нижнем левом направлении. Обратите внимание, что верхнее правое смещение равно координате x для всех диагоналей, начинающихся с левого края. Поправочный член (np.clip ...) предназначен для диагоналей, начинающихся с нижнего края. Аналогично, смещения в нижнем левом углу равны y, если мы начинаем с верхнего края, и нуждаемся в коррекции, если мы начинаем с правого края.

Answer 2

EDIT:

Вот версия в основном того же алгоритма, но без циклов:

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 есть последовательность (хотя это относится к общей последовательности для бесконечной сетки, в этом случае у вас будет одно перечисление, начинающееся в верхнем левом углу и еще один справа внизу).

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

Чтобы написать это, очень помогло нарисовать диаграмму.