Индексирование 2d массива с 3d-массивом в numpy

Question

У меня есть два массива.

"a", массив 2d numpy.

import numpy.random as npr 

a = array([[5,6,7,8,9],[10,11,12,14,15]]) 
array([[ 5, 6, 7, 8, 9], 
     [10, 11, 12, 14, 15]]) 

«IDX», 3d-NumPy массив, составляющий три варианта индекса, которые я хочу использовать для индекса «а».

idx = npr.randint(5, size=(nsamp,shape(a)[0], shape(a)[1])) 
array([[[1, 2, 1, 3, 4], 
     [2, 0, 2, 0, 1]], 

     [[0, 0, 3, 2, 0], 
     [1, 3, 2, 0, 3]], 

     [[2, 1, 0, 1, 4], 
     [1, 1, 0, 1, 0]]]) 

Теперь я хочу, чтобы индекс «а» три раза с индексами в «IDX», чтобы получить объект следующим образом:

array([[[6, 7, 6, 8, 9], 
     [12, 10, 12, 10, 11]], 

     [[5, 5, 8, 7, 5], 
     [11, 14, 12, 10, 14]], 

     [[7, 6, 5, 6, 9], 
     [11, 11, 10, 11, 10]]]) 

Наивный «а [IDX]» не работает. Любые идеи о том, как это сделать? (Я использую Python 3.4 и Numpy 1.9)

Answer 1

Вы можете использовать choose, чтобы сделать выбор из a:

>>> np.choose(idx, a.T[:,:,np.newaxis]) 
array([[[ 6, 7, 6, 8, 9], 
     [12, 10, 12, 10, 11]], 

     [[ 5, 5, 8, 7, 5], 
     [11, 14, 12, 10, 14]], 

     [[ 7, 6, 5, 6, 9], 
     [11, 11, 10, 11, 10]]]) 

Как вы можете видеть, a должен быть преобразован из массива с формой (2, 5) в массив с формой (5, 2, 1). Это по существу так, что он доступен для вещания с idx, который имеет форму (3, 2, 5).

(я узнал этот метод от @ immerrr Ответит здесь: https://stackoverflow.com/a/26225395/3923281)

Answer 2

Вы можете использовать метод take массива:

import numpy 

a = numpy.array([[5,6,7,8,9],[10,11,12,14,15]]) 

idx = numpy.random.randint(5, size=(3, a.shape[0], a.shape[1])) 

print a.take(idx)