时间戳

首页 > 解决方案 > 用numpy上的数组索引数组

python - 用numpy上的数组索引数组

问题描述

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

我有两个数组，
arr形状 (x, y, z)形状
索引(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，预期的数组将如下所示：

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 风格？

标签： pythonnumpy

解决方案

这可以使用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))

推荐阅读