У меня две матрицы, A и B:
A = array([[2., 13., 25., 1.], [ 18., 5., 1., 25.]])
B = array([[2, 1], [0, 3]])
Я хочу индексировать каждую строку A с каждой строкой B, создавая срез:
array([[25., 13.], [18., 25.]])
То есть, я действительно хочу что-то вроде:
array([A[i,b] for i,b in enumerate(B)])
Есть ли способ представить это прямо? Лучшее, что я могу сделать, это "плоский взлом":
A.flat[B + arange(0,A.size,A.shape[1])[:,None]]
@Ответ на Ophion велик и заслуживает внимания, но я хотел добавить некоторые объяснения и предложить более интуитивную конструкцию.
Вместо того, чтобы вращать B и затем возвращать результат назад, лучше просто повернуть arange. Я думаю, что это дает наиболее интуитивное решение, даже если оно принимает больше символов:
A[((0,),(1,)), B]
или эквивалентно
A[np.arange(2)[:, None], B]
Это работает, потому что вы действительно делаете массив i и массив j, каждый из которых имеет ту же форму, что и ваш желаемый результат.
i = np.array([[0, 0],
[1, 1]])
j = B
Но вы можете использовать только
i = np.array([[0],
[1]])
Потому что он будет транслироваться в соответствии с B (это то, что дает np.arange(2)[:,None]).
Наконец, чтобы сделать его более общим (не зная 2 как размер arange), вы также можете сгенерировать i из B с помощью
i = np.indices(B.shape)[0]
однако вы создаете i и j, вы просто называете это как
>>> A[i, j]
array([[ 25., 13.],
[ 18., 25.]])
Не очень, но:
A[np.arange(2),B.T].T
array([[ 25., 13.],
[ 18., 25.]])