Выберите один элемент в каждой строке массива с помощью столбцов

Есть ли лучший способ получить "output_array" из "input_array" и "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]

Ответы

Ответ 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])

Ответ 2

Я думаю, enumerate удобно.

[input_array[enum, item] for enum, item in enumerate(select_id)]

Ответ 3

Как насчет:

[input_array[x,y] for x,y in zip(range(len(input_array[:,0])),select_id)]

Ответ 4

(потому что я не могу опубликовать это как комментарий к принятому ответу)

Обратите внимание, что numpy.choose работает только в том случае, если у вас есть 32 или меньше вариантов (в этом случае размер вашего массива, по которому вы индексируете, должен иметь размер 32 или меньше). Кроме того, документация для numpy.choose говорит, что "Чтобы уменьшить вероятность неправильной интерпретации, даже если следующее" злоупотребление "номинально поддерживается, выбор не должны рассматриваться и не считаться единым массивом, т.е. самый внешний подобный последовательности контейнер должен быть либо списком, либо кортежем".

OP спрашивает:

  • Есть ли лучший способ получить "output_array" из "input_array" и "select_id"?
    • Я бы сказал, то, что вы изначально предлагали, кажется лучшим из представленных здесь. Это легко понять, масштабируется до больших массивов и эффективен.
  • Можно ли избавиться от диапазона (input_array.shape [0])?
    • Да, как показывают другие ответы, но принятый не работает вообще так хорошо, как то, что OP уже предлагает делать.