Могу ли я сделать разрез массива-пути как индексацию массива в Python?

Question

как я могу сделать нарезку массивов? Я хочу получить элемент 2,3; 7,8 в массиве a с учетом start_idx и end_idx. Я пытался a[[0,1], start_idx: end_idx], но не смог.

import numpy as np
a = np.array([[1,2,3,4], [5,6,7,8]])
start_idx=[1,2]
end_idx=[3,4]
a[[0,1], start_idx: end_idx]] <- failed
[a[0,1:3], a[1,2:4]] #[array([2, 3]), array([7, 8])] <- wanted result

Answer 1

Вы можете использовать понимание списка с enumerate и zip:

res = [a[idx, start: end] for idx, (start, end) in enumerate(zip(start_idx, end_idx))]

# [array([2, 3]), array([7, 8])]

Это похоже на решение @ Bazingaa, но итерирует по start_idx / end_idx напрямую, а не по позиционной индексации списка.

Повышение производительности незначительно и может проявиться только для больших массивов. Тем не менее, вы можете счесть перечисляющую логику более читабельной.

n = 100000
a = np.array([[1,2,3,4], [5,6,7,8]]*n)
start_idx = [1,2] * n
end_idx = [3,4] * n

def indexing(a, start_idx, end_idx):
    return [a[i, start_idx[i]:end_idx[i]] for i in range(len(start_idx))]

def enumerating(a, start_idx, end_idx):
    return [a[idx, start: end] for idx, (start, end) in enumerate(zip(start_idx, end_idx))]

%timeit indexing(a, start_idx, end_idx)     # 140 ms
%timeit enumerating(a, start_idx, end_idx)  # 124 ms

Answer 2

Вы можете использовать понимание списка:

result = [a[i, start_idx[i]:end_idx[i]] for i in range(len(start_idx))]
print (result) 

выход

[array([2, 3]), array([7, 8])]