如何根据包含初始索引和最终索引的 2 个一维数组对二维 numpy 数组进行切片

Question

我有一个 2D numpy 数组，假设它的形状为 4x10（4 行 10 列）。我有 2 个具有初始索引和最终索引的一维数组，因此它们都是 20x1。举个例子，假设

initial = [1, 2, 4, 5]
final = [3, 6, 8, 6]

那我想要

data[0,1:3]
data[1,2:6]
data[2,4:8]
data[3,5:6]

当然，这些数组中的每一个都会有不同的大小，所以我想将它们存储在一个列表中。如果我用 for 循环来做，它看起来像这样：

arrays = []
for i in range(4):
    slice = data[i,initial[i]:final[i]]
    arrays.append(slice)

有没有更有效的方法来做到这一点？我宁愿避免使用 for 循环，因为我的实际数据很大。

Answer 1

您可以使用 numpy.split with flattened data (using numpy.ndarray.flatten) 并修改切片：

sections = np.column_stack([initial, final]).flatten()

sections[::2] += np.arange(len(initial)) * data.shape[1]
sections[1::2] += sections[::2] - np.array(initial)

np.split(data.flatten(), sections)[1::2]

如何根据包含初始索引和最终索引的 2 个一维数组对二维 numpy 数组进行切片

how to slice 2D numpy array based on 2 1D arrays containing initial and final indexes

python

arrays

performance

numpy

slice