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

3 回答

• 0

  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 的前2个维度不匹配时，可以修改最后一种方法 .

  np.ix_ 可以简化此索引

  I, J = np.ix_(np.arange(4),np.arange(3))
  A[I, J, B]
  

  回复于 2024-05-22T10:56:02+08:00
• 0

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

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

  回复于 2024-05-22T10:56:02+08:00
• 2

  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) .

  回复于 2024-05-22T10:56:02+08:00