python - 矩阵逐行索引

Question

我有一个 numpy 数组，例如以下

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

还有另一个带有 bool 值的 numpy 数组，例如

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

我想获取其元素索引由 I 给出的矩阵。在上面的示例中，我想获取该矩阵

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

但如果我尝试

B = A[I]

然后我明白了

array([1, 5])

据我了解，这是因为每行中的 True 数量可能不同。但如果他们是呢？有没有办法使用 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]])