矩阵行式索引

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的数量在每一行中可能不相同。但是如果它们呢？有没有办法使用numpy这样做？

事实上，我想使用tensor模块在​​Theano中使用它。我有一个包含上述数组的上述（两个T.matrix theano变量）的theano表达式。有没有方便的方法来计算新的更小的矩阵？

Answer 1

如果您可以预先确定每行返回的项目数，则可以重新整形输出。我这样做：

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

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