矩阵行-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]])