NumPy Быстрый способ получить несколько строк по индексам

Question

Учитывая многомерный массив формы (6, 100, 2) и список индексов (например: [1,2,3,1,5,0,0,0,1]), и я хотел бы получить numpy ndarray с данными, как в этих индексах (имеется в виду, форма вывода равно (9, 100, 2), а строки 6,7,8 являются дубликатами друг друга).

Наивное решение:

arr = []
indexes = [1,2,3,1,5,0,0,0,1]
for i in indexes:
  arr.append(data[i])
arr = np.stack(arr)

Это наивное решение тратит много времени на __getitem__ и операция stack, и мне было интересно, есть ли более быстрый способ сделать это с numpy?

Answer 1

Вы можете просто использовать подписку. Таким образом, если data - ваша матрица 6 × 100 × 2, а indexes - ваш список индексов, вы можете получить результат с помощью:

result = data<b>[indexes]</b>