Векторизация этого цикла для Numpy

Question

Мне было интересно, как бы я векторизовал это для цикла.Учитывая 2x2x2 массив x и массив, где каждый элемент является ith, jth, and kth элементом массива, который я хочу получить x[i,j,k]

Учитывая массивы x и y

x = np.arange(8).reshape((2, 2, 2))
y = [[0, 1, 1], [1, 1, 0]]

Я хочу получить:

x[0, 1, 1] = 3 и x[1, 1, 0] = 6

Я пытался:

print(x[y]) 

Но он печатает:

array([[2, 3],
       [6, 7],
       [4, 5]])

Итак, я закончил:

for y_ in y:
    print(x[y_[0], y_[1], y_[2]])

Что работает, но я не могу не думать, что есть лучший способ.

Answer 1

Используйте транспонированный y т.е. zip(*y) в качестве индекса;У вас должны быть индексы для каждого измерения в качестве элемента для расширенного индексирования для работы:

x[tuple(zip(*y))]
# array([3, 6])