Чтобы проиллюстрировать и проверить вашу проблему, определите простые, реальные массивы:
In [44]: arr = np.arange(12).reshape(3,4)
In [45]: idx = np.array([[0,2,1,0],[2,2,1,0]])
In [46]: arr.shape
Out[46]: (3, 4)
In [47]: idx.shape
Out[47]: (2, 4)
Ваш желаемый расчет:
In [48]: res = np.zeros((2,4), int)
In [49]: for i in range(4):
...: res[:,i] = arr[:,i][idx[:,i]] # same as arr[idx[:,i], i]
...:
In [50]: res
Out[50]:
array([[0, 9, 6, 3],
[8, 9, 6, 3]])
Выполнение того же действия с одним шагом индексации:
In [51]: arr[idx, np.arange(4)]
Out[51]:
array([[0, 9, 6, 3],
[8, 9, 6, 3]])
Это передает два индексирующих массива друг на друга, а затем выбирает точки:
In [52]: np.broadcast_arrays(idx, np.arange(4))
Out[52]:
[array([[0, 2, 1, 0],
[2, 2, 1, 0]]),
array([[0, 1, 2, 3],
[0, 1, 2, 3]])]
Таким образом, мы индексируем массив (m, c) с 2 (n, c) массивами
То же самое:
arr [idx] arr [idx,:]
Используется idx
для выбора целых строк из arr
в результате получается форма idx
плюс последнее измерение arr
.Где то, что вы хотите, это просто элемент ith
строки idx[j,i]
.