索引張量與二進制矩陣在numpy

Question

我有張量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）。

Answer 1

可以使用列表理解做到這一點：

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

Answer 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)。

Answer 3

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

np.choose可如果最後的尺寸大於2使用（但小於32）。 （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 1日2種尺寸。

np.ix_馬y簡化此索引

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