NumPy: Есть ли способ выполнить «сдвиг» на ndarray?

Question

Исходный массив:

array([[ 8,  9, 10, 11],
       [12, 13, 14, 15]])

При смещении = 2 результат будет:

array([[ 10,  11, 12, 13],
       [14, 15, 0, 0]])

Со смещением = -2:

array([[ 0, 0, 8,  9],
       [10, 11, 12, 13]])

Есть ли способ сделать это, кроме создания нескольких фрагментов и объединения их с помощью np.r_?

Answer 1

Работает примерно так:

def shift1d(a, offset):
    if offset == 0:
        return a.copy()

    b = np.zeros_like(a)
    if offset<0:
        b.ravel()[-offset:] = a.ravel()[:offset]
    else:
        b.ravel()[:-offset] = a.ravel()[offset:]

    return b

shift1d(a, 2)
# array([[10, 11, 12, 13],
#        [14, 15,  0,  0]])

shift1d(a,-2)
# array([[ 0,  0,  8,  9],
#        [10, 11, 12, 13]])