我有一个3D数组a
数据和一个2D数组b
索引。我需要使用a
中的索引沿第3轴采用b
的子数组。我可以使用take
这样做:
a = np.arange(24).reshape((2,3,4))
b = np.array([0,2,1,3]).reshape((2,2))
np.array([np.take(a_,b_,axis=1) for (a_,b_) in zip(a,b)])
我可以在没有列表理解的情况下使用一些花哨的索引吗?我担心效率,所以如果在这种情况下花式索引不是更有效,我想知道它。
编辑我尝试的第一件事是a[[0,1],:,b]
,但它没有提供我需要的子数组
答案 0 :(得分:2)
In [317]: a
Out[317]:
array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
In [318]: a = np.arange(24).reshape((2,3,4))
...: b = np.array([0,2,1,3]).reshape((2,2))
...: np.array([np.take(a_,b_,axis=1) for (a_,b_) in zip(a,b)])
...:
Out[318]:
array([[[ 0, 2],
[ 4, 6],
[ 8, 10]],
[[13, 15],
[17, 19],
[21, 23]]])
所以你想要0&第1块的2列,1& 3从第二个。
制作一个与c
形状匹配的b
,并体现这种观察
In [319]: c=np.array([[0,0],[1,1]])
In [320]: c
Out[320]:
array([[0, 0],
[1, 1]])
In [321]: b
Out[321]:
array([[0, 2],
[1, 3]])
In [322]: a[c,:,b]
Out[322]:
array([[[ 0, 4, 8],
[ 2, 6, 10]],
[[13, 17, 21],
[15, 19, 23]]])
这是正确的数字,但形状不正确。
可以使用列向量代替c
。
In [323]: a[np.arange(2)[:,None],:,b] # or a[[[0],[1]],:,b]
Out[323]:
array([[[ 0, 4, 8],
[ 2, 6, 10]],
[[13, 17, 21],
[15, 19, 23]]])
至于形状,我们可以转置最后两个轴
In [324]: a[np.arange(2)[:,None],:,b].transpose(0,2,1)
Out[324]:
array([[[ 0, 2],
[ 4, 6],
[ 8, 10]],
[[13, 15],
[17, 19],
[21, 23]]])
这个转置是必需的,因为我们在两个索引数组之间有一个切片,它是基本索引和高级索引的混合。这是有记录的,但从来没有那么令人费解。它将切片尺寸(3)放在最后,我们必须将其转置回来。
很好的小索引拼图!
此高级/基本转置的最新问题和解释:
Indexing numpy multidimensional arrays depends on a slicing method
答案 1 :(得分:1)
这是我的第一次尝试。我会看看能不能做得更好。
#using numpy broadcasting.
np.r_[a[0][:,b[0]],a[1][:,b[1]]].reshape(2,3,2)
Out[300]: In [301]:
array([[[ 0, 2],
[ 4, 6],
[ 8, 10]],
[[13, 15],
[17, 19],
[21, 23]]])
第二次尝试:
#convert both a and b to a 2d array and then slice all rows and only columns determined by b.
a.reshape(6,4)[np.arange(6)[:,None],b.repeat(3,0)].reshape(2,3,2)
Out[429]:
array([[[ 0, 2],
[ 4, 6],
[ 8, 10]],
[[13, 15],
[17, 19],
[21, 23]]])