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

Question

I am trying to insert a 2D array of size [2, 2] into a 3D array of size [2, 3, 2]. For every page of the 3D array (axis=0), the position to insert the 2D array (read: row number) may be different. I tried to use the np.insert function. However, I am struggling...

import numpy as np

arr = np.arange(12).reshape(2, 3, 2)

arr
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]]])

row_number_before_insertion = [1, 2]
val_to_insert = (np.ones(4) * 100).reshape(2,2)
arr_expanded = np.insert(arr, row_number_before_insertion , val_to_insert, axis=1)

arr_expanded
array([[[  0,   1],
        [100, 100],
        [  2,   3],
        [100, 100],
        [  4,   5]],

       [[  6,   7],
        [100, 100],
        [  8,   9],
        [100, 100],
        [ 10,  11]]])

I am actually looking for the following result:

arr_expanded
array([[[  0,   1],
        [100, 100],
        [100, 100],
        [  2,   3],
        [  4,   5]],

       [[  6,   7],
        [  8,   9],
        [100, 100],
        [100, 100],
        [ 10,  11]]])

Answer 1

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_expanded



arr = 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_expanded
array([[[   0.,    1.],
        [ 100.,  100.],
        [ 100.,  100.],
        [   2.,    3.],
        [   4.,    5.]],

       [[   6.,    7.],
        [   8.,    9.],
        [ 100.,  100.],
        [ 100.,  100.],
        [  10.,   11.]]])

Answer 2

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

from skimage.util.shape import view_as_windows

def 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]: arr
Out[44]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]]])

In [45]: row_number_before_insertion
Out[45]: array([1, 2])

In [46]: val_to_insert
Out[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