在 Numpy 中使用 2d 数组对 3d 数组进行切片

Question

我有一个关于使用 2d 数组切片 3d 数组的问题。 largearray 是我想用 2d smallarray 中的值切片的 3d 数组

    array([[[36.914   , 38.795   , 37.733   , 36.68    , 35.411003,
     33.494   , 36.968002, 39.902   , 43.943   , 48.398   ],
    [37.121   , 38.723   , 37.706   , 36.653   , 35.491997,
     33.638   , 36.697998, 39.668   , 43.817   , 48.551   ]],

   [[37.292   , 28.454   , 23.414   , 23.018   , 21.83    ,
     19.472   , 28.364   , 35.492   , 28.786999, 36.23    ],
    [37.04    , 28.256   , 23.135   , 22.937   , 21.839   ,
     19.382   , 28.517   , 35.816   , 28.922   , 36.509   ]]],

largearray.shape = (2, 2, 10)

smallarray = array([[5, 7],[9, 3]]) smallarray.shape = (2, 2)

3d 数组的结果应该被切分直到对应的 2d 数组的值。结果应如下所示：

    array([[[36.914   , 38.795   , 37.733   , 36.68    , 35.411003],
   [37.121   , 38.723   , 37.706   , 36.653   , 35.491997, 33.638   ,
   36.697998]],
   [[37.292   , 28.454   , 23.414   , 23.018   , 21.83    , 19.472   ,
   28.364   , 35.492   , 28.786999],
   [37.04 , 28.256, 23.135]]])

最终的计算将在非常大的数组上进行，因此如果计算的计算成本尽可能低，那就太好了。

希望你能帮我解决这个问题！

Answer 1

如果将 3D 数组和 2D 数组分别转换为 2D 数组和 1D 数组，计算会容易一些。

largearray = largearray.reshape(-1,largearray.shape[-1])
smallarray = smallarray.reshape(-1)

ans = np.array([largearray[i,:smallarray[i]].tolist() for i in range(len(smallarray))]).reshape(2,2)