python - 从具有索引列表的多维数组中选择(selecting from a multi-dimesional array with a list of indices)

Question

Let's say I have an array with size batch x max_len x output_size , where batch , max_len , and output_size all correspond to positive natural numbers.

(假设我有一个大小为batch x max_len x output_size ，其中batch ， max_len和output_size都对应于正自然数。)

I have a list of indices which correspond to individual items in dimension 1 (ie max_len ).

(我有一个索引列表，它对应于维度1中的各个项目（即max_len ）。)

How can I select from the array given these indices?

(给定这些索引，如何从数组中选择？)

As a concrete example, let's say I have the following:

(举一个具体的例子，假设我有以下内容：)

>>> l = np.random.randn(4,5,6)
>>> l.shape
(4, 5, 6)
>>> idx = [0,0,2,3]

When I select l given idx I get:

(当我选择l给idx我得到：)

>>> l[:,idx,:].shape
(4, 4, 6)
>>>

I also tried np.take but reached the same result:

(我也尝试了np.take但达到了相同的结果：)

>>> np.take(l,idx,axis=1).shape
(4, 4, 6)
>>> 

However, the output I am looking after is (4,1,6) as I am trying to have only one item looking at each element in the batch (ie first dimension).

(但是，我要处理的输出是(4,1,6)因为我试图只让一个项目查看batch中的每个元素（即第一维）。)

How can I produce the output with the proper shape?

(如何产生形状合适的输出？)

  ask by Clement Attlee translate from so

Answer 1

Use np.take_along_axis after extending idx to have same ndims as l -

(扩展idx后使用np.take_along_axis以具有与l相同的ndims-)

np.take_along_axis(l,np.asarray(idx)[:,None,None],axis=1)

With explicit integer-array indexing -

(使用显式整数数组索引-)

l[np.arange(len(idx)),idx][:,None] # skip [:,None] for (4,6) shaped o/p