Python：沿不同的行将2D数组插入3D NumPy数组

我正在尝试将大小为

[2,2]

的二维数组插入大小为

[2,3,2]

的三维数组。对于三维阵列的每一页（轴=0），插入二维阵列（读取：行号）的位置可能不同。我尝试使用

np.insert

函数。然而，我正在挣扎

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]]])

我实际上在寻找以下结果：

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

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

---

这是一个基于数组分配和

掩蔽的-

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

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

样本运行-

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]]])

另一个具有
重复
和
掩蔽
-

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

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

下面是一个使用
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.]]])

发布的解决方案对你有用吗？