如何在多维 NumPy 数组的轴中选择两个数组之一，该轴具有一维索引数组

Question

我有一个形状为 (n, 2, 3) 的数组：

array = np.array([[[-0.903, -3.47, -0.946], [-0.883, -3.48, -0.947]],
                  [[-1.02, -3.45,  -0.992], [-1.01,  -3.46,     -1]],
                  [[-1.02, -3.45,  -0.992], [-0.998, -3.45,     -1]],
                  [[-0.638, -3.5,  -0.897], [-0.604, -3.51, -0.896]],
                  [[-0.596, -3.52, -0.896], [-0.604, -3.51, -0.896]]])

和第二个轴的索引数组，其中每个值引用两个组合中的每一个，例如对于[-0.903, -3.47, -0.946], [-0.883, -3.48, -0.947]如果索引数组中对应的值为1，则必须取[-0.883, -3.48, -0.947]：

indices = np.array([0, 1, 0, 0, 1], dtype=np.int64)

结果数组必须如下所示，形状为 (n, 3):

[-0.903, -3.47, -0.946] [-1.01, -3.46, -1] [-1.02, -3.45, -0.992] [-0.638, -3.5, -0.897] [-0.604, -3.51, -0.896]

仅通过 NumPy.

我怎么能在指定的维度上做到这一点

Answer 1

在 numpy 中，您可以沿两个维度组合切片。如果您执行 arr[idx_x, idx_y]，其中 idx_x 和 idx_y 是相同长度的一维数组，您将获得元素数组：[arr[idx_x[0], idx_y[0]], arr[idx_x[1], idx_y[1]], arr[idx_x[2], idx_y[2]], ...]

在您的示例中，如果您这样做：

indices = np.array([0, 1, 0, 0, 1], dtype=np.int64)
x_idxs = np.arange(len(indices), dtype=int)
print(array[x_idxs, indices])

这将 return 您想要的结果。

Answer 2

尝试使用 for 循环：

out = []
for i in range(len(indices)):
    out.append(list(arr[i,indices[i]]))

print(out)

Output:
[[-0.903, -3.47, -0.946],
 [-1.01, -3.46, -1.0],
 [-1.02, -3.45, -0.992],
 [-0.638, -3.5, -0.897],
 [-0.604, -3.51, -0.896]]

如何在多维 NumPy 数组的轴中选择两个数组之一，该轴具有一维索引数组

How to pick one of two arrays in an axis of multidimensional NumPy array with an 1D index array for that axis

python

numpy

indexing