矩陣行索引

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做到這一點嗎？

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

Answer 1

如果您可以預先弄清楚每行返回了多少項，則只需調整輸出的形狀即可。 我會這樣：

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

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