Numpy: многомерная индексация?

Question

Я странствую, можно ли индексировать сразу несколько измерений? С некоторым вещанием. Пример:

Предположим, у меня есть массив A в форме (n,d). Предположим, у меня есть индексный массив, скажем, I с целочисленными значениями от 0 до d-1. Установите B = A[:,I].

Если shape(I) == (k,), для k, то B имеет форму (n,k) и B[x,y] = A[x,I[y]].

Но если shape(I) == (k,p) для любых (k,p), тогда я хотел, чтобы B имел форму (n,k,p) с B[x,y,z] = A[x,I[y,z]].

1 ° Как я могу получить это поведение?

2 ° Есть ли у него недостаток, которого я не видел?

Answer 1

Вы можете сделать это точно так, как вы это описали:

import numpy as np
n = 100
d = 20
k = 10
p = 17

A = np.random.random((n, d))
I = np.random.randint(low=0, high=d, size=(k, p))
B = A[:, I]
print(B.shape)  # (n, k, p)

# Testing if the new array B is constructed as expected
x = 3
y = 5
z = 7
print(B[x, y, z])
print(A[x, I[y, z]])
print(B[x, y, z] == A[x, I[y, z]])

Трудно ответить, является ли это хорошей реализацией или нет, без контекста. Но, в общем, неплохо использовать numpy и векторизацию, если вы думаете о скорости.