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

我有两个数组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]])

目前我的做法很好:

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

我的笔记本电脑平均用了674ms(8核,16 ram)