python - 如何根据包含初始和最终索引的 2 个 1D 数组对 2D numpy 数组进行切片

Question

我有一个 2D numpy 数组，假设它的形状为 4x10（4 行和 10 列）。我有 2 个具有初始和最终索引的 1D 数组，所以它们都是 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循环，因为我的实际数据很大。

标签： pythonarraysperformancenumpyslice

Answer 1

您可以使用numpy.split扁平化data（使用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]