Как отмечается в комментариях, вы можете просто индексировать в массив arr
, используя массив idx
.
In [47]: arr
Out[47]: array([6, 7, 8, 9])
In [48]: idx
Out[48]:
array([[3, 2, 2, 0],
[0, 3, 2, 3],
[3, 2, 2, 3]])
In [49]: arr[idx]
Out[49]:
array([[9, 8, 8, 6],
[6, 9, 8, 9],
[9, 8, 8, 9]])
Если вы хотите подход, который является менее волшебным и более информативным, тогда ниже будет более полезным.
# flatten the `idx` array; index into `arr`; then reshape to `idx's` shape.
In [50]: arr[idx.ravel()].reshape(idx.shape)
Out[50]:
array([[9, 8, 8, 6],
[6, 9, 8, 9],
[9, 8, 8, 9]])