Как дополнить 2d массив с переменной длиной на строку слева и справа, чтобы сформировать больший 2d массив

Question

У меня есть 2d массив, как

small = np.arange(9).reshape((3, 3))

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

Я хочу применить заполнение каждой строки для заполнения переменной 0s слева. И чтобы убедиться, что полученный массив 2d имеет форму 3 x 8. (заполнение нулями справа)

offset = np.array([1, 3, 2])

чтобы результат выглядел как

array([[ 0.,  0.,  1.,  2.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  3.,  4.,  5.,  0.,  0.],
       [ 0.,  0.,  6.,  7.,  8.,  0.,  0.,  0.]])

Какой лучший способ добиться этого?

благодаря решению @Divakar. Я провел несколько тестов по следующим методам.

def f1(small, offset, ncols):
    nrows, num_small_cols = small.shape
    big = np.zeros((nrows, ncols))
    inner = np.empty_like(small, dtype=np.int64)
    for i in range(num_small_cols):
        inner[:, i] = offset + i
    big[np.arange(nrows)[:, None], inner] = small
    return big

def f2(small, offset, ncols):
    n = small.shape[1]
    r = np.arange(ncols)
    offset2 = offset[:,None]
    # This took a lot of time
    mask = (offset2 <= r) & (offset2 + n > r)
    out = np.zeros_like(mask, dtype=np.float64)
    out[mask] = small.ravel()
    return out

def f3(small, offset, ncols):
    n = small.shape[1]
    m = ncols - n
    small_pad = np.zeros((len(small), n + 2*m))
    small_pad[:,m:m+n] = small    
    w = view_as_windows(small_pad, (1,ncols))[:,:,0]
    return w[np.arange(len(offset)), ncols-offset-n]

n = 10000
offset = np.repeat(np.array([1, 3, 2]), n)
small = np.random.rand(n * 3, 5)

%timeit f1(small, offset, 9)
# 1.32 ms

%timeit f2(small, offset, 9)
# 2.24 ms

%timeit f3(small, offset, 9)
# 1.3 ms

Answer 1

Подход № 1

Мы можем использовать broadcasting, чтобы создать маску для присвоения этим позициям и затем присвоить массиву, инициализированному нулями -

def pad_offsetpos(small, ncols):
    n = small.shape[1]
    r = np.arange(ncols)
    mask = (offset[:,None] <= r) & (offset[:,None]+n > r)
    out = np.zeros(mask.shape)
    out[mask] = small.ravel()
    return out

Пробный прогон -

In [327]: small
Out[327]: 
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

In [328]: offset
Out[328]: array([1, 3, 2])

In [329]: pad_offsetpos(small, ncols=8)
Out[329]: 
array([[0., 0., 1., 2., 0., 0., 0., 0.],
       [0., 0., 0., 3., 4., 5., 0., 0.],
       [0., 0., 6., 7., 8., 0., 0., 0.]])

Подход № 2

Мы также можем использовать np.lib.stride_tricks.as_strided на основе scikit-image's view_as_windows для эффективного извлечения патчей после заполнения входного массива достаточным количеством нулей с обеих сторон -

from skimage.util.shape import view_as_windows

def pad_offsetpos_strided(small, ncols):
    n = small.shape[1]
    m = ncols - n
    small_pad = np.zeros((len(small), n + 2*m))
    small_pad[:,m:m+n] = small    
    w = view_as_windows(small_pad, (1,ncols))[:,:,0]
    return w[np.arange(len(offset)), ncols-offset-n]

Answer 2

PaddedMat=numpy.zeros(shape=[3,8],dtype="float")

и затем цикл, чтобы заполнить его.

PaddedMat[i,offset[i]:offset[i]+3]=small[i,:]

и т.д ...