如何使用另一个 3d 数组索引 3d 数组 (t, x, y)答案

【问题标题】：How to index 3d array (t, x, y) using another 3d array如何使用另一个 3d 数组索引 3d 数组 (t, x, y)
【发布时间】：2023-03-17 23:39:01
【问题描述】：

我有一个 3d 值数组，

vals = np.array([
    [
        [10, 20, 30],
        [40, 50, 60],
    ],
    [
        [15, 25, 35],
        [45, 55, 65],
    ],
])

以及相应的 3d 坐标数组

coords = np.array([
    [
        [0,1],
        [0,2],
        [1,1]
    ],
    [
        [0,0],
        [1,1],
        [1,2]
    ]
])

coords 的每个最内层数组表示对应于 vals 中的二维数组之一的 (x,y) 坐标。例如coords中的坐标[0,1]对应的值为20，coords中的坐标[1,2]对应的值为65。

如何以这种方式使用coords 对vals 进行子集化？

我可以像这样解决这个具体的例子

np.array([
    vals[0][coords[0][:, 0], coords[0][:, 1]],
    vals[1][coords[1][:, 0], coords[1][:, 1]]
])

array([[20, 30, 50],
       [15, 55, 65]])

但显然我想要一个更动态的解决方案。

【问题讨论】：

高级索引：vals[np.arange(len(coords))[:,None], coords[..., 0], coords[..., 1]]

标签： python arrays numpy

【解决方案1】：

有趣的是，写下我的问题似乎总是让我找到答案。盯着答案矩阵，

array([[20, 30, 50],
       [15, 55, 65]])

我问自己，“我如何从原始索引值中重现这个矩阵？”。例如，要提取值 20，我知道我可以做到

vals[0, 0, 1]

如果我想提取答案中的第一行值，[20, 30, 50] 我应该这样做

vals[[0,0,0], [0,0,1], [1,2,1]]

那么要得到完整的答案矩阵，我应该这样做

vals[[[0,0,0],[1,1,1]], [[0,0,1],[0,1,1]], [[1,2,1],[0,1,2]]]

从这里开始，我将重点放在生成这三个索引矩阵上。它们可以按如下方式构造：

i1 = np.arange(coords.shape[0])[:, None].repeat(coords.shape[1], axis=1)
i2 = coords[:,:,0]
i3 = coords[:,:,1]

# Thus the generalized solution
vals[i1, i2, i3]

这个答案与@Psidom 在 cmets 中提到的高级索引解决方案极为相似，但可能不太优雅。

【讨论】：