按列索引在Numpy数组的每一行中选择一个元素

Question

有没有更好的方法从"input_array"和"select_id"获取"output_array"？

我们可以摆脱range( input_array.shape[0] )？

>>> input_array = numpy.array( [ [3,14], [12, 5], [75, 50] ] )
>>> select_id = [0, 1, 1]
>>> print input_array
[[ 3 14]
 [12  5]
 [75 50]]

>>> output_array = input_array[  range( input_array.shape[0] ), select_id ]
>>> print output_array
[ 3  5 50]

Answer 1

您可以从给定的数组中进行选择,使用numpy.choose该数组构造索引数组(在您的情况下select_id)和一组数组(在您的情况下input_array)中的数组可供选择.但是,您可能首先需要转置input_array以匹配尺寸.以下是一个小例子:

In [101]: input_array
Out[101]: 
array([[ 3, 14],
       [12,  5],
       [75, 50]])

In [102]: input_array.shape
Out[102]: (3, 2)

In [103]: select_id
Out[103]: [0, 1, 1]

In [104]: output_array = np.choose(select_id, input_array.T)

In [105]: output_array
Out[105]: array([ 3,  5, 50])

Answer 2

(因为我不能将此作为对已接受答案的评论发布)

请注意,numpy.choose仅当您有32个或更少的选项时才有效(在这种情况下,您要编制索引的数组的维度必须大小为32或更小).此外,该文件的numpy.choose说

为了减少误解的可能性,即使名义上支持以下"滥用",选择既不应该也不应被视为单个数组,即,最外层的序列式容器应该是列表或元组.

OP问:

1. 有没有更好的方式来获得output_array从input_array和select_id？
   • 我会说,你最初建议的方式似乎是这里提出的最好的方式.它易于理解,可扩展到大型数组,并且效率很高.
2. 我们可以摆脱range(input_array.shape[0])？
   • 是的,正如其他答案所示,但是被接受的答案并不像OP已经建议做的那么好.