在numpy中用二进制矩阵索引张量答案

【问题标题】：Indexing tensor with binary matrix in numpy在numpy中用二进制矩阵索引张量
【发布时间】：2015-05-14 08:43:43
【问题描述】：

我有一个张量 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)。

【问题讨论】：

标签： python numpy matrix

【解决方案1】：

np.where 来救援。和mtrw's原理一样的回答：

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]

【讨论】：

【解决方案2】：

您可以使用列表理解来做到这一点：

C = np.array([[A[i, j, B[i, j]] for j in range(A.shape[1])] for i in range(A.shape[0])])

【讨论】：

【解决方案3】：

C = A[:,:,0]*(B==0) + A[:,:,1]*(B==1) 应该可以工作。如果您需要索引更多平面，您可以将其概括为 np.sum([A[:,:,k]*(B==k) for k in np.arange(A.shape[-1])], axis=0)。

【讨论】：