在numpy

Question

有没有办法避免使用for循环并通过使用一些索引调用arr来获得结果？潜在的dim1等于50 000，dim2等于1000，dim3等于3。

import numpy as np
dim1 = 10
dim2 = 2
dim3 = 3
arr = np.arange(60).reshape(dim1,dim2,dim3)
arr2 = np.arange(dim1*dim2).reshape(dim1,dim2)
np.mod(arr2,dim3,out=arr2)

res = []
rng = np.arange(dim1)
for x in range(dim2):
    sl = arr2[:,x]
    temp = arr[rng,x,sl]
    res.append(temp)
res = np.asarray(res).T

基本上，我想从arr中提取值，这是一个3D数组，但矩阵arr2指示要选择的列。

最佳