Question

我需要能够返回多维数组的一部分，但我不知道如何以正确的方式执行此操作。我这样做的方式似乎很天真：

import numpy as np
a=np.ones([3,3,3,3,3])
b=np.asarray([2,2])
c=np.asarray([2,2])
print a[b[0],b[1],:,c[0],c[1]]

并将返回

[1,1,1]

然而我想要的是这样的：

a=np.ones([3,3,3,3,3])
b=np.asarray([2,2])
c=np.asarray([2,2])
print a[b,:,c]

返回a本身，虽然我希望它返回[1,1,1]。

我不知道为什么。如何在不指定逐个元素的情况下读取数组的一部分，但是将我想要的数组的索引作为包提供？

P.S。感谢@hcwhsa，我更新了问题以更具体地解决我想要的问题。

Answer 1

将b定义为元组：

>>> b = (2, 2)
>>> a[b]
array([ 1.,  1.,  1.])

或者在将其传递给tuple

之前将其转换为a[]

>>> b = np.asarray([2,2])
>>> a[tuple(b)]
array([ 1.,  1.,  1.])

Answer 2

我可以想到两种方法来做到这一点，也不是完美的。一种是滚动你想要完成的轴：

ax = 2 # the axis you want to have all values in
np.rollaxis(a, ax, a.ndim)[tuple(np.r_[b,c])]

如果你将两个轴向后移动，这适用于a[b,:,:,c]（在轴编号的索引位移时要小心！）

np.rollaxis(np.rollaxis(a, ax, a.ndim), ax, a.ndim)[tuple(np.r_[b,c])]

其中np.rollaxis(a, ax, a.ndim)移动轴ax，您希望保留所有结尾：

a = np.zeros((1,2,3,4,5))
a.shape
#(1,2,3,4,5)
np.rollaxis(a, ax, a.ndim).shape
#(1,2,4,5,3)

np.r_[b,c]只是连接两个数组。你也可以这样做：tuple(np.concatenate([b,c]))

或者，你可以使用我评论中的那个：

a[tuple(b) + (slice(None),) + tuple(c)]

其中slice是start:end:step语法创建的对象。 None为您提供:，但您可以动态创建它（无需在正确的位置键入:）。因此，a[1:3]相当于a[slice(1,3)]，a[:3]是a[slice(None,3)]。我把它包装在元组中，以便可以“添加”到其他两个元组中以创建一个长元组。

如何在Python中返回多维数组的子数组？

2 个答案: