如何在 Numpy 中使用 3d 索引數組進行二維數組切片

Question

我有 2 arrays 作為輸入。 在數組上為 output。數組a保存數據並且形狀為(N,M) ，而數組b保存索引並且形狀為(N,X,2) 。 結果數組的形狀應為(N,X) ，其值取自a 。

現在它只適用於 for 循環。 由於我有巨大的 arrays 作為輸入，我該如何對其進行矢量化？

下面是一個示例代碼，用於演示我現在擁有的內容：

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)

Answer 1

好的，所以我對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)

我不確定它是否是最佳解決方案，但它確實有效。