Использование двух массивов индексации в массиве numpy

Question

У меня есть 3-мерный массив Numpy, и я пытаюсь проиндексировать его по второму и третьему измерениям, используя два одномерных массива NumPy.То есть:

np.random.seed(0)
dims = (1,5,5)
test_array = np.arange(np.prod(dims)).reshape(*dims)
test_array

производит:

array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24]]])

, если я затем создаю два массива для его индексации:

idx_dim1 = np.array([0,2,4])
idx_dim2 = np.array([1,3])

Я считаю, что не могупримените их оба одновременно:

test_array[:,idx_dim1,idx_dim2]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-193-95b23ed3210c> in <module>()
----> 1 test_array[:,idx_dim1,idx_dim2]

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,) (2,) 

Я могу разделить их, скажем,

test_array[:, idx_dim1,:][:,:,idx_dim2]

, что дает

array([[[ 1,  3],
        [11, 13],
        [21, 23]]])

Но это работает тольков смысле только для чтения - я не могу присвоить значения test_array следующим образом, так как я назначаю срез, а не оригинал.Я также не совсем уверен, почему совместное использование двух индексаторов не работает.Есть ли какой-то концептуальный нюанс, который я здесь не понимаю?И есть ли хорошее решение, которое позволяет мне присваивать значения?

Спасибо!

Answer 1

Попробуйте

test_array[:, idx_dim1[:,None], idx_dim2] = some_value

Где,

>>> test_array[:, idx_dim1[:,None], idx_dim2]
array([[[ 1,  3],
        [11, 13],
        [21, 23]]])