使用 numpy 上的数组索引数组

Question

它类似于围绕 SO 的一些问题，但我不太了解获得我想要的东西的技巧。

我有两个arrays，
形状为 (x, y, z) 的arr
形状 (x, y) 的索引，其中包含对 z 感兴趣的索引。

对于索引的每个值，我想在arr中获取实际值，其中：

arr.x == indexes.x  
arr.y == indexes.y  
arr.z == indexes[x,y]

这将给出一个类似于索引形状的 shape(x,y) 数组。

例如：

arr = np.arange(99)
arr = arr.reshape(3,3,11)
indexes = np.asarray([
[0,2,2],
[1,2,3],
[3,2,10]])
# indexes.shape == (3,3)

# Example for the first element to be computed
first_element = arr[0,0,indexes[0,0]]

使用上述indexes ，预期的 arrays 将如下所示：

expected_result = np.asarray([
[0,13,24],
[34,46,58],
[69,79,98]])

我试过elements = np.take(arr, indexes, axis=z)但它给出了一个形状数组 (x, y, x, y)
我也尝试过elements = arr[indexes, indexes,:]之类的东西，但我没有得到我想要的。
我看到了一些涉及转置索引并将其转换为元组的答案，但我不明白它会有什么帮助。

注意：我对 numpy 有点陌生，所以我还没有完全理解索引。
你会如何解决这个 numpy 风格？

Answer 1

这可以使用np.take_along_axis来完成

import numpy as np

#sample data
np.random.seed(0)
arr = np.arange(3*4*2).reshape(3, 4, 2) # 3d array
idx = np.random.randint(0, 2, (3, 4))   # array of indices

out = np.squeeze(np.take_along_axis(arr, idx[..., np.newaxis], axis=-1))

在这段代码中，索引数组又增加了一个轴，因此它可以广播到我们从中进行选择的数组arr的形状。 然后，由于np.take_along_axis的返回值与索引数组的形状相同，我们需要使用np.squeeze删除这个额外的维度。

另一种选择是使用np.choose ，但在这种情况下，您进行选择所沿的轴必须移动到数组的第一个轴：

out = np.choose(idx, np.moveaxis(arr, -1, 0))

Answer 2

这里的解决方案应该适合你： Indexing 3d numpy array with 2d array

适应您的代码：

ax_0 = np.arange(arr.shape[0])[:,None]
ax_1 = np.arange(arr.shape[1])[None,:]

new_array = arr[ax_0, ax_1, indexes]

Answer 3

您可以使用np.take_along_axis执行此类操作，该操作只能沿一维应用，因此您需要重塑输入和索引。

您要执行的操作是：

out[i, j] = arr[i, j, indices[i, j]]

然而，我们被迫将arr和indices重塑，即map (i, j)到k ，这样我们就可以应用np.take_along_axis 。 将进行以下操作：

out[k] = arr[k, indices[k]] # indexing along axis=1

---

这里的实际用法归结为：

>>> put = np.take_along_axis(arr.reshape(9, 11), indices.reshape(9, 1), axis=1)
array([[ 0],
       [13],
       [24],
       [34],
       [46],
       [58],
       [69],
       [79],
       [91]])

然后重塑回indices的形状：

>>> put.reshape(indices.shape)
array([[ 0, 13, 24],
       [34, 46, 58],
       [69, 79, 91]])