首页 文章

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

提问于
浏览
0

我有一个三维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有一些概念上的细微差别吗?是否有一个很好的解决方案,允许我分配值?

谢谢!

1 回答

  • 1

    尝试

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

    哪里,

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

相关问题