Получение значений из местоположений на основе трехмерной матрицы в двумерной матрице

Question

Предполагая, что у меня есть матрица numpy: data = np.random.rand(200, 50, 100), и у меня были местоположения, для которых мне нужны значения: locs = np.random.randint(50, size=(200, 2)).

Как получить полученную матрицу формы (200, 2, 100)? По сути, я хотел бы получить значения из data в местах, указанных locs.

Если я сделаю: data[locs], я получу в результате полученную матрицу формы (200, 2, 50, 100), а не (200, 2, 100).

Обновлена ​​с более подробной информацией в соответствии с запросом:

Если у нас есть:

data = np.arange(125)
reshaped = np.reshape(data, (5, 5, 5))
locs = [[3, 4], [2, 1], [1, 3], [3, 3], [0, 0]]

Тогда выполнение чего-то вроде data[locs] должно дать следующий вывод:

array([[[ 15,  16,  17,  18,  19],
        [ 20,  21,  22,  23,  24]],

       [[ 35,  36,  37,  38,  39],
        [ 30,  31,  32,  33,  34]],

       [[ 55,  56,  57,  58,  59],
        [ 65,  66,  67,  68,  69]],

       [[ 90,  91,  92,  93,  94],
        [ 90,  91,  92,  93,  94]],

       [[100, 101, 102, 103, 104],
        [100, 101, 102, 103, 104]]])

Answer 1

Результатом расширенного индекса будет форма индексов по измерениям, которые вы индексируете. data[locs] эквивалентно data[locs, :, :], поэтому ваша фигура будет locs.shape + data.shape[1:] или (200, 2, 50, 100).

То, что вы представляете для запроса, это индексировать ось 1 of data с использованием locs, удерживая ось 0 в шаге со строкой в ​​locs. Для этого вам нужно индексировать с помощью locs вдоль оси 1 и указывать индекс от 0 до 200 по оси ноль.

Важно помнить, что все расширенные индексы должны транслироваться на такой же формы Поскольку locs имеет форму (200, 2), первый индекс должен иметь форму (200, 1) или (200, 2) для правильной трансляции. Я покажу первое, так как оно проще и эффективнее.

data = np.random.rand(200, 50, 100)
locs = np.random.randint(50, size=(200, 2))
rows = np.arange(200).reshape(-1, 1)

result = data[rows, locs, :]