Python：在3D numpy数组中随机插入多行

我有一个3D数组

arr

，大小

（2,5,5）

。我还有另一个数组

行到行

，大小

（3,5）

我想在

arr

的每页中随机插入

行。但是，
行到插入
不能作为块插入。此外，插入位置应为任意页面的
arr

然而，我正在努力有效地将
行插入到
中。我当前的解决方案包含一个for循环

import numpy as np

arr = np.arange(100, 125).reshape(5, 5)
arr = np.repeat(arr[None, :, :], 2, axis=0)

rows_to_ins = np.random.randint(0, 99, (3,5))

row_nums_3D = np.random.randint(0, arr.shape[1], (2, 1, 3))

arr_ins = list()
for i in range(row_nums_3D.shape[0]):
    arr_ins.append(np.insert(arr[i, :, :], np.squeeze(row_nums_3D[i, :, :]), rows_to_ins, axis=0))
arr_ins = np.asarray(arr_ins)

我想知道，我是否可以避免for循环。矢量化解决方案是什么样的

也许一个更具体的例子会有助于理解我的问题

# arr - shape (2, 5, 5)

[[[100 101 102 103 104]
  [105 106 107 108 109]
  [110 111 112 113 114]
  [115 116 117 118 119]
  [120 121 122 123 124]]

 [[100 101 102 103 104]
  [105 106 107 108 109]
  [110 111 112 113 114]
  [115 116 117 118 119]
  [120 121 122 123 124]]]


# rows_to_insert - shape(3, 5) 

[[37 31 28 34 10]
 [ 2 97 89 36 11]
 [66 14 70 37 45]]

我正在寻找这样一个潜在的结果：

# 3D array with insertet rows - shape (2, 8, 5)

[[[100 101 102 103 104]
  [ 2 97 89 36 11]
  [66 14 70 37 45]
  [105 106 107 108 109]
  [110 111 112 113 114]
  [115 116 117 118 119]
  [120 121 122 123 124]
  [37 31 28 34 10]]

 [[66 14 70 37 45]
  [100 101 102 103 104]
  [105 106 107 108 109]
  [ 2 97 89 36 11]
  [110 111 112 113 114]
  [37 31 28 34 10]
  [115 116 117 118 119]
  [120 121 122 123 124]]]

这是一种矢量化方法-

def insert_random_places(arr, rows_to_ins):
    m,n,r = arr.shape
    N = len(rows_to_ins) + n
    idx = np.random.rand(m,N).argsort(1)
    out = np.zeros((m,N,r),dtype=np.result_type(arr, rows_to_ins))
    np.put_along_axis(out,np.sort(idx[:,:n,None],axis=1),arr,axis=1)
    np.put_along_axis(out,idx[:,n:,None],rows_to_ins,axis=1)
    return out

样本运行-

In [58]: arr
Out[58]: 
array([[[100, 101, 102, 103, 104],
        [105, 106, 107, 108, 109],
        [110, 111, 112, 113, 114],
        [115, 116, 117, 118, 119]],

       [[100, 101, 102, 103, 104],
        [105, 106, 107, 108, 109],
        [110, 111, 112, 113, 114],
        [115, 116, 117, 118, 119]]])

In [59]: rows_to_ins
Out[59]: 
array([[77, 72,  9, 20, 80],
       [69, 79, 47, 64, 82]])

In [60]: np.random.seed(0)

In [61]: insert_random_places(arr, rows_to_ins)
Out[61]: 
array([[[100, 101, 102, 103, 104],
        [ 69,  79,  47,  64,  82],
        [105, 106, 107, 108, 109],
        [110, 111, 112, 113, 114],
        [115, 116, 117, 118, 119],
        [ 77,  72,   9,  20,  80]],

       [[100, 101, 102, 103, 104],
        [ 77,  72,   9,  20,  80],
        [ 69,  79,  47,  64,  82],
        [105, 106, 107, 108, 109],
        [110, 111, 112, 113, 114],
        [115, 116, 117, 118, 119]]])

另一个基于
掩蔽
-

def insert_random_places_v2(arr, rows_to_ins):
    m,n,r = arr.shape
    L = len(rows_to_ins)
    N = L + n

    insert_idx = np.random.rand(m,N).argpartition(kth=-L,axis=1)[:,-L:]
    mask = np.zeros((m,N),dtype=bool)
    np.put_along_axis(mask,insert_idx,1,axis=1)

    out = np.zeros((m,N,r),dtype=np.result_type(arr, rows_to_ins))
    rows_to_ins_3D = rows_to_ins[np.random.rand(m,L).argsort(1)]
    out[mask] = rows_to_ins_3D.reshape(-1,r)
    out[~mask] = arr.reshape(-1,r)
    return out

您能验证您的
掩蔽
解决方案吗<代码>输出
以前没有定义过。@谢谢您的反馈。编辑。