картирование через 2 numpy.ndarray одновременно

Question

Вот и проблема. Скажем, у меня есть матрица А =

array([[ 1., 0., 2.], 
     [ 0., 0., 2.], 
     [ 0., -1., 3.]]) 

и вектор индексов р = array([0, 2, 1]). Я хочу превратить матрицу 3x3 A в массив длины 3 (назовите его v), где v [j] = A [j, p [j]] для j = 0, 1, 2. Я могу сделать это следующим образом :

v = map(lambda (row, idx): row[idx], zip(A, p)) 

Таким образом, для указанных выше матриц а и вектора индексов р я ожидаю получить array([1, 2, -1]) (т.е. 0-й элемент строки 0, второй элемент строки 1, первый элемент строки 2).

Но могу ли я достичь того же результата, используя собственный numpy (т. Е. Без явного перекоса и затем отображения)? Благодарю.

Answer 1

Я не думаю, что такая функциональность существует. Чтобы достичь того, чего вы хотите, я могу думать о двух простых способах. Вы можете сделать:

np.diag(A[:, p]) 

Здесь массив p применяется как индекс столбца для каждой строки таким образом, что по диагонали вы будете иметь элементы, которые вы ищете.

В качестве альтернативы вы можете избежать производить много ненужных записей с помощью:

A[np.arange(A.shape[0]), p]