Нарезка массива NumPy по заданным начальным и конечным индексам для общих измерений

Question

Дан простой массив x формы (N_1...N_k), где k произвольно, и 2 массива:

start_indices=[a_1,...,a_k], end_indices=[b_1,...b_k], where `0<=a_i<b_i<=N_i`.

Я хочу нарезать х следующим образом: x[a_1:b_1,...,a_k:b_k].

Скажем:

x is of shape `(1000, 1000, 1000)`
start_indices=[450,0,400]
end_indices=[550,1000,600].

Я хочу, чтобы вывод был равен x[450:550,0:1000,400:600].

Например, я попытался определить:

slice_arrays = (np.arange(start_indices[i], end_indices[i]) for i in range(k))

и используйте

x[slice_arrays]

но это не сработало.

Answer 1

Вы можете использовать нотацию slice для создания кортежа индексации, который можно использовать для индексации -

indexer = tuple([slice(i,j) for (i,j) in zip(start_indices,end_indices)])
out = x[indexer]

В качестве альтернативы, с сокращением np.s_ -

indexer = tuple([np.s_[i:j] for (i,j) in zip(start_indices,end_indices)])

или с map для компактного -

indexer = tuple(map(slice,start_indices,end_indices))