Python - 在每一行的不同位置对数组进行切片

我有一个 2D python 数组，我想以一种奇怪的方式对其进行切片 - 我想要一个恒定宽度的切片，从每行的不同位置开始。如果可能的话，我想以矢量化的方式做到这一点。

例如我有数组A=np.array([range(5), range(5)])看起来像

array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])

我想按如下方式对其进行切片：每行 2 个元素，从位置 0 和 3 开始。起始位置存储在b=np.array([0,3])。期望的输出是：np.array([[0,1],[3,4]]) i.e.

array([[0, 1],
       [3, 4]])

我试图得到这个结果的明显事情是A[:,b:b+2]但这不起作用，而且我找不到任何可以的东西。

速度很重要，因为这将在循环中对较大的数组进行操作，并且我不想成为代码其他部分的瓶颈。

您可以使用np.take():

In [21]: slices = np.dstack([b, b+1])

In [22]: np.take(arr, slices)
Out[22]: 
array([[[0, 1],
        [3, 4]]])