我有一个 3D 数组和一个 2D 索引数组。如何选择最后一个轴?
import numpy as np
# example array
shape = (4,3,2)
x = np.random.uniform(0,1, shape)
# indices
idx = np.random.randint(0,shape[-1], shape[:-1])
这是一个可以给出所需结果的循环。但应该有一种有效的矢量化方法来做到这一点。
result = np.zeros(shape[:-1])
for i in range(shape[0]):
for j in range(shape[1]):
result[i,j] = x[i,j,idx[i,j]]
一个可能的解决方案:
np.take_along_axis(x, np.expand_dims(idx, axis=-1), axis=-1).squeeze(axis=-1)
或者,
i, j = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
x[i, j, idx]
针对2D的修正,首先使用meshgrid构建笛卡尔映射。
m=np.meshgrid(range(shape[0]), range(shape[1]), indexing="ij")
results = x[m[0], m[1], idx]