使用2d数组索引3d numpy数组

时间:2016-10-24 10:50:36

标签: python arrays numpy

我正在尝试从3d numpy数组中提取值。目前我可以执行以下操作:

newmesh.shape
(40,40,40)

newmesh[2,5,6]
6

但是,如果我尝试使用数组对其进行索引,则结果不符合预期;

newmesh[np.array([2,5,6])].shape
(3, 42, 42)

我尝试过使用np.take,但它产生以下内容;

np.take(newmesh,np.array([2,5,6]))
[-1 -1 -1]

为什么会发生这种情况?我的目标是输入一个(n,3)数组,其中每一行对应一个newmesh的值,即输入一个(n,3)数组将返回一个长度为n的1d数组。

1 个答案:

答案 0 :(得分:1)

使用idx作为(n,3)索引数组,使用linear-indexing的一种方法将使用np.ravel_multi_index -

np.take(newmesh,np.ravel_multi_index(idx.T,newmesh.shape))

元组形成的方法看起来像这样 -

newmesh[tuple(idx.T)]

如果只有三个维度,您甚至可以使用柱状切片来索引每个维度,如下所示 -

newmesh[idx[:,0],idx[:,1],idx[:,2]]

运行时测试如果有人有兴趣查看与列出的方法相关的性能数字,那么这是一个快速的运行时测试 -

In [18]: newmesh = np.random.rand(40,40,40)

In [19]: idx = np.random.randint(0,40,(1000,3))

In [20]: %timeit np.take(newmesh,np.ravel_multi_index(idx.T,newmesh.shape))
10000 loops, best of 3: 22.5 µs per loop

In [21]: %timeit newmesh[tuple(idx.T)]
10000 loops, best of 3: 20.9 µs per loop

In [22]: %timeit newmesh[idx[:,0],idx[:,1],idx[:,2]]
100000 loops, best of 3: 17.2 µs per loop