Python: Insert 2D array into 3D NumPy array along different rows

Here's one based on array-assignment and masking -

from skimage.util.shape import view_as_windowsdef insert_into_arr(arr, row_number_before_insertion, val_to_insert):    ma,na,ra = arr.shape    L = len(val_to_insert)    N = len(row_number_before_insertion)    out = np.zeros((ma,na+L,ra),dtype=arr.dtype)    mask = np.ones(out.shape, dtype=bool)    w = view_as_windows(out,(1,L,1))[...,0,:,0]    w[np.arange(N), row_number_before_insertion] = val_to_insert.T    wm = view_as_windows(mask,(1,L,1))[...,0,:,0]    wm[np.arange(N), row_number_before_insertion] = 0    out[mask] = arr.ravel()    return out

Sample run -

In [44]: arrOut[44]: array([[[ 0,  1],        [ 2,  3],        [ 4,  5]],       [[ 6,  7],        [ 8,  9],        [10, 11]]])In [45]: row_number_before_insertionOut[45]: array([1, 2])In [46]: val_to_insertOut[46]: array([[784, 659],       [729, 292],       [935, 863]])In [47]: insert_into_arr(arr, row_number_before_insertion, val_to_insert)Out[47]: array([[[  0,   1],        [784, 659],        [729, 292],        [935, 863],        [  2,   3],        [  4,   5]],       [[  6,   7],        [  8,   9],        [784, 659],        [729, 292],        [935, 863],        [ 10,  11]]])

Another with repeat and masking -

def insert_into_arr_v2(arr, row_number_before_insertion, val_to_insert):      ma,na,ra = arr.shape    r = row_number_before_insertion    L = len(val_to_insert)    M = na+L    out = np.zeros((ma,na+L,ra),dtype=arr.dtype)    idx = ((r + M*np.arange(len(r)))[:,None] + np.arange(L)).ravel()    out.reshape(-1,ra)[idx] =np.repeat(val_to_insert[None],ma,axis=0).reshape(-1,ra)    mask = np.isin(np.arange(ma*(na+L)),idx, invert=True)    out.reshape(-1,ra)[mask] = arr.reshape(-1,ra)    return out

Here's a solution using vstack:

def insert_into_arr(arr, row_number_before_insertion, val_to_insert):    num_slices, num_rows, num_cols = arr.shape    arr_expanded = np.zeros((num_slices, num_rows + val_to_insert.shape[0], num_cols))    for i in range(num_slices):        if row_number_before_insertion[i] == 0:            arr_expanded[i, :, :] = np.vstack((val_to_insert, arr[i, :, :]))        else:            arr_expanded[i, :, :] = np.vstack((arr[i, 0:row_number_before_insertion[i], :], val_to_insert, arr[i, row_number_before_insertion [i]:, :]))    return arr_expandedarr = np.arange(12).reshape(2, 3, 2)row_number_before_insertion = [1, 2]val_to_insert = (np.ones(4) * 100).reshape(2,2)arr_expanded = insert_into_arr(arr, row_number_before_insertion, val_to_insert)arr_expandedarray([[[   0.,    1.],        [ 100.,  100.],        [ 100.,  100.],        [   2.,    3.],        [   4.,    5.]],       [[   6.,    7.],        [   8.,    9.],        [ 100.,  100.],        [ 100.,  100.],        [  10.,   11.]]])