Рулонная матрица со всеми возможными сменами

Question

У меня есть функция, которая принимает матрицу (h, w) и создает тензор порядка 3, который для каждой строки содержит все возможные сдвиги размера c этой строки.Вот функция и пример использования1012 * с c = 2 дает матрицу

1 2
2 3
3 0

, и это происходит для каждой строки, в результате чего получается тензор.

Мой вопрос: как сделать это быстрее? Полагаю, в идеале я хотел бы избавиться от цикла for, но любое более быстрое решение приветствуется.

Answer 1

Вы можете использовать stride_tricks.

def fast_roll(v, c):
    *h, w = v.shape
    V = np.zeros((*h, w+c-1), v.dtype)
    V[..., :w] = v
    return np.lib.stride_tricks.as_strided(V, (*h, w, c), (*V.strides, V.strides[-1]))

Обратите внимание, что это создает несмежное представление.При необходимости сделайте непрерывную копию.

Пример:

>>> fast_roll(np.arange(9).reshape(3, 3), 2)
array([[[0, 1],
        [1, 2],
        [2, 0]],

       [[3, 4],
        [4, 5],
        [5, 0]],

       [[6, 7],
        [7, 8],
        [8, 0]]])