python - 在 numpy 数组上使用两个索引数组

Question

我有一个三维 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 分配值，因为我分配给切片而不是原始切片。我也不太清楚为什么同时使用两个索引器不起作用。我在这里不理解 numpy 的一些概念上的细微差别吗？有没有一个好的解决方案可以让我分配值？

谢谢!

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]]])