Выбор конкретного столбца в каждой строке из массива

Question

Я пытаюсь выбрать конкретные элементы столбца для каждой строки массива. Например, в следующем примере:

In [1]: a = np.random.random((3,2))
Out[1]: 
array([[ 0.75670668,  0.1283942 ],
       [ 0.51326555,  0.59378083],
       [ 0.03219789,  0.53612603]])

Я хотел бы выбрать первый элемент первого ряда, второй элемент второго ряда и первый элемент третьего ряда. Поэтому я попытался сделать следующее:

In [2]: b = np.array([0,1,0])

In [3]: a[:,b]

Но это приводит к следующему выводу:

Out[3]: 
array([[ 0.75670668,  0.1283942 ,  0.75670668],
       [ 0.51326555,  0.59378083,  0.51326555],
       [ 0.03219789,  0.53612603,  0.03219789]])

что явно не то, что я ищу. Есть ли простой способ сделать то, что я хотел бы сделать без использования циклов?

Answer 1

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

a[np.arange(3), (0,1,0)]

в вашем примере выше.

Answer 2

ОК, просто для пояснения, давайте сделаем простой пример

A=diag(arange(0,10,1))

дает

array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 2, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 3, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 4, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 6, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 7, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 8, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 9]])

тогда

A[0][0:4]

1010 * дает *

array([0, 0, 0, 0])

это первая строка, элементы от 0 до 3. Но

A[0:4][1]

не дает первые 4 строки, 2-й элемент в каждой. Вместо этого мы получаем

array([0, 1, 0, 0, 0, 0, 0, 0, 0, 0])

то есть весь 2-й столбец.

A[0:4,1]

дает

array([0, 1, 0, 0])

Я уверен, что для этого есть очень веская причина, которая имеет смысл для программистов. но для тех из нас, кто не имеет отношения к этой великой религии, это может сбить с толку.

Answer 3

Это не столько ответ, сколько попытка немного документировать это. Для ответа выше у нас будет:

>>> import numpy as np
>>> A = np.array(range(6))
>>> A
array([0, 1, 2, 3, 4, 5])
>>> A.shape = (3,2)
>>> A
array([[0, 1],
       [2, 3],
       [4, 5]])
>>> A[(0,1,2),(0,1,0)]
array([0, 3, 4])

Указание списка (или кортежа) отдельных координат строки и столбца позволяет необычное индексирование массива. Первый пример в комментарии на первый взгляд выглядит аналогично, но индексы - это кусочки. Они не распространяются на весь диапазон, и форма возвращаемого массива отличается:

>>> A[0:2,0:2]
array([[0, 1],
       [2, 3]])

Для второго примера в комментарии

>>> A[[0,1],[0,1]]
array([0, 3])

Таким образом, похоже, что срезы разные, но за исключением того, что независимо от того, как составлены индексы, вы можете указать кортеж или список (значения x, значения y) и восстановить эти конкретные элементы из массива.