按列索引在Numpy数组的每一行中选择一个元素 [英] Select One Element in Each Row of a Numpy Array by Column Indices

问题描述

是否有更好的方法从"input_array"和"select_id"中获取"output_array"?

Is there a better way to get the "output_array" from the "input_array" and "select_id" ?

我们可以摆脱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]

推荐答案

您可以使用

You can choose from given array using numpy.choose which constructs an array from an index array (in your case select_id) and a set of arrays (in your case input_array) to choose from. However you may first need to transpose input_array to match dimensions. The following shows a small example:

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])

这篇关于按列索引在Numpy数组的每一行中选择一个元素的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！