Разделить массив в Python, используя индексы из списка

Question

У меня есть двумерный массив размером 3 на 7 в numpy:

[[1 2 3 4 5 6 7]
[4 5 6 7 8 9 0]  
[2 3 4 5 6 7 8]]

У меня также есть список, который содержит индексы точек разделения:

[1, 3]

Теперь я хочу разделить массив, используя индексы в списке, чтобы я получил:

[[1 2]
[4 5]
[2 3]]

[[ 2 3 4]
[5 6 7]
[3 4 5]]

[[ 4 5 6 7]
[7 8 9 0]
[5 6 7 8]]

Как я могу сделать это в Python?

Answer 1

Вы можете использовать понимание списка с нарезкой, используя zip для парного извлечения индексов.

A = np.array([[1, 2, 3, 4, 5, 6, 7],
              [4, 5, 6, 7, 8, 9, 0],
              [2, 3, 4, 5, 6, 7, 8]])

idx = [1, 3]
idx = [0] + idx + [A.shape[1]]

res = [A[:, start: end+1] for start, end in zip(idx, idx[1:])]

print(*res, sep='\n'*2)

[[1 2]
 [4 5]
 [2 3]]

[[2 3 4]
 [5 6 7]
 [3 4 5]]

[[4 5 6 7]
 [7 8 9 0]
 [5 6 7 8]]