У меня есть трехмерный исходный массив numpy и двумерный массив numpy индексов.
Например:
src = np.array([[[1,2,3],[4,5,6]],
[[7,8,9],[10,11,12]]])
idx = np.array([[0,1],
[1,2]])
Я хотел бы получить 2d-массив, где каждый элемент представляет индексированное значение в самом внутреннем измерении в этой позиции:
array([[1,5],
[8,12]])
Как мне сделать это с помощью numpy?
Вы можете попробовать np.take, вот документация < / а>.
Однако вы должны подсчитать индекс массива после сглаживания всех элементов. Например, вы должны использовать
src = np.array([[[1,2,3],[4,5,6]],
[[7,8,9],[10,11,12]]])
idx = np.array([[0,4],
[7,11]])
# Wanted result
res = np.take(src, idx)
где src считалось [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
Вы также можете попробовать np.take_along_axis, вот документация.
При использовании этого метода ваши src и idx должны быть одинакового размера, поэтому сначала вы должны разжать src и сжать res.
# Unsqueezed the last dim
idx = np.expand_dims(idx, axis=-1)
# Squeeze the last dim
res = np.take_along_axis(src, idx, axis=2).squeeze(-1)
Вы можете использовать метод np.choose с небольшим изменением формы :
np.choose(idx.reshape((1, 2, 2)), src.transpose()).reshape((2, 2))
>>>> array([[ 1, 8],
[ 5, 12]])
Прямая индексация:
src[np.arange(2)[:, None], np.arange(2), idx]
