Python Numpy从多维数组中选择由索引矩阵指定的矩阵
我有一个小数组Python Numpy从多维数组中选择由索引矩阵指定的矩阵,python,arrays,numpy,indexing,Python,Arrays,Numpy,Indexing,我有一个小数组a,大小5x5x4x5x5。我有另一个矩阵b,大小5x5。我想将I的a[I,j,b[I,j]]从0到4,以及j从0到4。这将给我一个5x5x1x5x5矩阵。有没有办法不使用2进行循环?让我们把矩阵a想象成100(=5x5x4)大小的矩阵(5,5)。所以,如果你能为每个三元组得到一个线性索引-(i,j,b[i,j]),你就完成了。这就是np.ravel\u multi\u index的用武之地。下面是代码 import numpy as np import itertools #
a
,大小5x5x4x5x5
。我有另一个矩阵b
,大小5x5
。我想将I
的a[I,j,b[I,j]]
从0到4,以及j
从0到4。这将给我一个5x5x1x5x5
矩阵。有没有办法不使用2进行循环?让我们把矩阵a
想象成100(=5x5x4)
大小的矩阵(5,5)
。所以,如果你能为每个三元组得到一个线性索引-(i,j,b[i,j])
,你就完成了。这就是np.ravel\u multi\u index
的用武之地。下面是代码
import numpy as np
import itertools
# create some matrices
a = np.random.randint(0, 10, (5, 5, 4, 5, 5))
b = np.random(0, 4, (5, 5))
# creating all possible triplets - (ind1, ind2, ind3)
inds = list(itertools.product(range(5), range(5)))
(ind1, ind2), ind3 = zip(*inds), b.flatten()
allInds = np.array([ind1, ind2, ind3])
linearInds = np.ravel_multi_index(allInds, (5,5,4))
# reshaping the input array
a_reshaped = np.reshape(a, (100, 5, 5))
# selecting the appropriate indices
res1 = a_reshaped[linearInds, :, :]
# reshaping back into desired shape
res1 = np.reshape(res1, (5, 5, 1, 5, 5))
# verifying with the brute force method
res2 = np.empty((5, 5, 1, 5, 5))
for i in range(5):
for j in range(5):
res2[i, j, 0] = a[i, j, b[i, j], :, :]
print np.all(res1 == res2) # should print True
正是为了这个目的-
np.take_along_axis(a,b[:,:,None,None,None],axis=2)