Создайте массив 3d numpy, чьи записи увеличиваются в заданном направлении

Question

Учитывая форму (n, m, k) и ненулевой вектор (u, v, w) с записями в {-1,0,1}, я хотел бы создать массив numpy формы (n , т, к). Записи массива должны начинаться с 1 и увеличиваться по направлению к вектору.

Хотя я специально запрашиваю трехмерные массивы, позвольте мне проиллюстрировать их на двухмерных примерах:

(n,m) = (3,4) and (u,v) = (-1,0) gives:

4 3 2 1
4 3 2 1
4 3 2 1

(n,m) = (4,3) and (u,v) = (1,-1) gives:

1 2 3
2 3 4
3 4 5
4 5 6

I Я могу создать их, используя вложенные циклы for, но мне интересно, есть ли более быстрое решение, так как я буду работать с большими массивами.

Answer 1

Вот метод, использующий stride_tricks. Работает для произвольного числа измерений.

from numpy.lib.stride_tricks import as_strided                    

def pp(dims,strides):                                                      
    dims,strides = np.asarray(dims),np.asarray(strides)                    
    aux = np.arange(1,(dims-1).sum()+2)                                    
    return as_strided(aux[(dims-1)@(strides==-1):],dims,aux.strides*strides)

Примеры:

>>> pp((2,3),(-1,0))
array([[2, 2, 2],
       [1, 1, 1]])
>>> pp((2,3,4),(-1,-1,-1))
array([[[7, 6, 5, 4],
        [6, 5, 4, 3],
        [5, 4, 3, 2]],

       [[6, 5, 4, 3],
        [5, 4, 3, 2],
        [4, 3, 2, 1]]])

Обратите внимание, что для оси y принято начинать сверху вниз и опускаться вниз. Если вы хотите иначе, вам придется перевернуть его.

Также обратите внимание, что созданные массивы являются несмежными представлениями, если вы хотите изменить их, лучше сделайте копию.