如何在 Numpy 中使用 3d 索引数组进行 2d 数组切片答案

【问题标题】：How can I use a 3d array of indices for a 2d array slicing in Numpy如何在 Numpy 中使用 3d 索引数组进行 2d 数组切片
【发布时间】：2022-10-04 20:23:39
【问题描述】：

我有 2 个数组作为输入。在数组上作为输出。数组a 保存数据，形状为(N,M)，而数组b 保存索引，形状为(N,X,2)。结果数组的形状应为(N,X)，其值取自a。

现在它只适用于 for 循环。既然我有巨大的数组作为输入，我怎么能矢量化它？

下面是一个示例代码来演示我现在所拥有的：

import numpy as np

# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)

a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))

t_result = np.empty((10, 9))

for i in range(b.shape[0]):
    t_result[i] = a[i, b[i, :, 0]]

print(t_result)
print(t_result.shape)

【问题讨论】：

标签： python numpy vectorization numpy-slicing

【解决方案1】：

好的，所以我对另一个post from scleronomic 的答案做了一些调整：

import numpy as np

# a of shape (N,M)
# b of shape (N,X,2)
# t_result of shape (N, X)

a = np.random.randint(0, 10, size=(10, 10))
b = np.random.randint(0, 2, size=(10, 9, 2))

t_result = np.empty((10, 9))

t_result = a[np.arange(a.shape[0])[:,None],b[np.arange(b.shape[0]),:,0]]

print(t_result)
print(t_result.shape)

我不确定它是否是最好的解决方案，但它确实有效。

【讨论】：