在最后一个轴上选择numpy数组

2nc8po8w 于 2023-08-05 发布在其他

关注(0)|答案(2)|浏览(87)

我有一个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]]

型

2条答案

可能的解决方案：

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建立笛卡尔Map。

m=np.meshgrid(range(shape[0]), range(shape[1]), indexing="ij")
results = x[m[0], m[1], idx]

字符串