我有一个张量 A 使得 A.shape = (32, 19, 2) 和一个二进制矩阵 B 使得 B.shape = (32, 19)。是否有我可以执行的单行操作来获得矩阵 C,其中 C.shape = (32, 19) 和 C(i,j) = A[i, j, B[i,j]]?
基本上,我想使用 B 作为索引矩阵,如果 B[i,j] = 1,我将 A[i,j,1] 形成 C(i,j)。
最佳答案
np.where
来拯救。原理和mtrw的
回答是一样的:
In [344]: A=np.arange(4*3*2).reshape(4,3,2)
In [345]: B=np.zeros((4,3),dtype=int)
In [346]: B[[0,1,1,2,3],[0,0,1,2,2]]=1
In [347]: B
Out[347]:
array([[1, 0, 0],
[1, 1, 0],
[0, 0, 1],
[0, 0, 1]])
In [348]: np.where(B,A[:,:,1],A[:,:,0])
Out[348]:
array([[ 1, 2, 4],
[ 7, 9, 10],
[12, 14, 17],
[18, 20, 23]])
如果最后一个维度大于 2(但小于 32),则可以使用 np.choose
。 (choose
在列表或第一个维度上操作,因此 rollaxis
。
In [360]: np.choose(B,np.rollaxis(A,2))
Out[360]:
array([[ 1, 2, 4],
[ 7, 9, 10],
[12, 14, 17],
[18, 20, 23]])
B
也可以直接作为索引使用。诀窍是以广播到相同形状的方式指定其他维度。
In [373]: A[np.arange(A.shape[0])[:,None], np.arange(A.shape[1])[None,:], B]
Out[373]:
array([[ 1, 2, 4],
[ 7, 9, 10],
[12, 14, 17],
[18, 20, 23]])
当 B
与 A
的第一个和第二个维度不匹配时,可以修改最后一种方法以工作。
np.ix_
可以简化这个索引
I, J = np.ix_(np.arange(4),np.arange(3))
A[I, J, B]
关于python - 在 numpy 中用二进制矩阵索引张量,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/30225539/