矩陣行-Wise索引

Question

我有一個numpy的陣列，例如，下面的

import numpy as np 
A = np.array([[1, 2, 3], [4, 5, 6]]) 

，並與布爾值另numpy的陣列，例如，

I = np.array([[True, False, False], [False, True, False]]) 

我想獲得的矩陣其元素的索引由I給出。在上面的例子中，我想得到矩陣

array([[1], [5]]) 

但是如果我嘗試

B = A[I] 

然後我得到

array([1, 5]) 

我明白，這是由於事實的Trues「號s可能不是每行中的相同。但是，如果他們是？有沒有使用numpy這樣做的方法？

事實上，我想在Theano中使用這個，使用tensor模塊。我有一個包含上述數組的theano表達式（兩個T.matrix theano變量）。有沒有方便的方式來計算新的更小的矩陣？

Answer 1

如果您可以預先確定每行返回多少項目，則可以重新設計輸出。我會這樣做：

n = I.sum(1).max() 
x = A[I].reshape(-1, n) 

print(x) 
array([[1], 
     [5]])