python - 在 numpy 中从 3d 获取特定的 2d 数组

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

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

最好的

标签： pythonarraysnumpyvectorizationarray-broadcasting