In [129]: a = [np.array([[1, 2, 3],
...: [4, 5, 6],
...: [7, 8, 9]]),
...: np.array([[11, 12, 13],
...: [14, 15, 16],
...: [17, 18, 19]]),
...: np.array([[99, 98, 97],
...: [96, 95, 94],
...: [93, 92, 91]])]
In [130]: b = np.array([[0, 0, 1],
...: [2, 1, 0],
...: [2, 1, 2]])
In [131]:
In [131]: A = np.array(a)
In [132]: A.shape
Out[132]: (3, 3, 3)
Вы хотите использовать b
для индексации первого измерения. Для других измерений вам нужны индексы, которые транслируются с b
, то есть вектор столбца и вектор строки:
In [133]: A[b, np.arange(3)[:,None], np.arange(3)]
Out[133]:
array([[ 1, 2, 13],
[96, 15, 6],
[93, 18, 91]])
Существуют различные вспомогательные функции для создания этих массивов, например,
In [134]: np.ix_(range(3),range(3))
Out[134]:
(array([[0],
[1],
[2]]), array([[0, 1, 2]]))
и ogrid
, как указано в другом ответе.
Вот относительно новая функция, которая также выполняет свою работу:
In [138]: np.take_along_axis(A, b[None,:,:], axis=0)
Out[138]:
array([[[ 1, 2, 13],
[96, 15, 6],
[93, 18, 91]]])
Мне пришлось немного подумать, прежде чем я получил правильную настройку b
.