如何解决如何在特定轴上索引由索引给定的多维数组?
假设我有一个形状为A
的4d数组(D0,D1,D2,D3)
。我有一个形状为B
的一维数组(D0,)
,其中包含我在第2轴处需要的索引。
实现我所需的简单方法:
output_lis = []
for i in range(D0):
output_lis.append(A[i,:,B[i],:])
#output = np.concatenate(output_lis,axis=0) #it is wrong to use concatenate. Thanks to @Mad Physicist. Instead,using stack.
output = np.stack(output_lis,axis=0) #shape: [D0,D3]
所以,我的问题是如何用numpy API快速实现它?
解决方法
使用花式索引在锁步中沿二维移动。在这种情况下,arange
提供序列i
,而B
提供序列B[i]
:
A[np.arange(D0),:,B,:]
与(D0,D1,D3)
循环结果的形状不同,此数组的形状确实为for
。
要从示例中获得相同的结果,请使用stack
(添加新轴),而不要使用concatenate
(使用现有轴):
output = np.stack(output_lis,axis=0)
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。