Python Numpy - 3-мерные индексы в 2-мерном массиве без петли

Question

Я хочу построить массив V1 формы (n,p,q) с использованием массива индексов idx той же формы, примененного к массиву V0 формы (p,q).Способ его построения с помощью цикла следующий:

for i in range(n):
    V1[i,:,:] = V0[idx[i,:,:],range(q)]

Другими словами, массив idx[i,:,:] содержит индексы для элементов 1-го измерения V0.Я применяю его со связанным индексом 2-го измерения, захваченного в range(q), чтобы получить соответствующий элемент по первому измерению окончательного массива V1.

Мой вопрос заключается в следующем: есть лиспособ построения V1 без зацикливания, с использованием методов вещания / индексации?

Спасибо.

Answer 1

Как насчет

V1 = V0[idx, range(q)] #?

Пример:

import numpy as np

# set up dummy data
n,p,q = 3,4,5
V1 = np.empty((n,p,q))
V0 = np.random.rand(p,q)
idx = np.random.randint(0,n,(n,p,q))

# original
V1_old = V1.copy()
for i in range(n):
    V1_old[i,:,:] = V0[idx[i,:,:],range(q)]

# new
V1_new = V0[idx, range(q)]

# test
print(np.array_equal(V1_old, V1_new)) # True