Numpy：使用2D索引数组从数组中提取行的快速方法

Question

我有2个数组a和b：

N,D,V,W = 2,3,4,5
a = np.random.randint(0,V,N*D).reshape(N,D)
a
array([[2, 3, 3],
       [2, 0, 3]])

b = np.random.randint(0,10,V*W).reshape(V,W)
b
array([[0, 1, 0, 5, 5],
       [0, 3, 6, 8, 7],
       [8, 8, 9, 0, 9],
       [4, 6, 3, 3, 1]])

我需要做的是使用数组a的元素值作为数组b的行索引，用数组b的行替换数组a的每个元素。

此刻我正在以这种方式运行，效果很好：

b[a.ravel(),:].reshape(*a.shape,-1)
array([[[8, 8, 9, 0, 9],
        [4, 6, 3, 3, 1],
        [4, 6, 3, 3, 1]],

       [[8, 8, 9, 0, 9],
        [0, 1, 0, 5, 5],
        [4, 6, 3, 3, 1]]])

但是，这种方法似乎有点慢。

我用以下方法进行了测试：

N,D,V,W = 20000,64,100,256

我的笔记本电脑（8核，16 ram）平均花费674ms

有人可以推荐一种更快但仍然简单的方法吗？