如何通過 2d boolean 數組的行有效地索引 numpy 1d arrays

Question

我希望能夠通過 2d boolean 數組的行來索引 1d 數組。 我知道有方法可以為 1d boolean arrays 執行此操作，但由於效率對我來說很重要，所以我不知道有什么方法不僅for循環。 一個例子：

我有一個二維掩碼 (Nxd) 和一個一維數組 (d,)，我想從中按行索引：

mask = [[False, True, False, True], 
        [False, True, True, False]]

y = [0, 1, 2, 3]

從上面，我希望得到：

y_masked = [[1, 3]
            [1, 2]]

我試過使用np.where來索引 boolean arrays，但我無法將 1d 數組轉換回正確的 2d 數組，而且，我發現的結果形狀不正確。 我也試過為每個i簡單地計算y[mask[i]] ，但這很慢。 我的主要問題是無法找到非逐行方法。

Answer 1

In [29]: mask = np.array([[False, True, False, True],
    ...:         [False, True, True, False]])
    ...: 
    ...: y = np.array([0, 1, 2, 3])
In [30]: 
In [30]: mask
Out[30]: 
array([[False,  True, False,  True],
       [False,  True,  True, False]])
In [31]: y
Out[31]: array([0, 1, 2, 3])

首先是明顯的逐行屏蔽：

In [32]: [y[row] for row in mask]
Out[32]: [array([1, 3]), array([1, 2])]

如果我們創建一個在形狀上與mask匹配的數組，我們會得到：

In [33]: Y = y[None,:].repeat(2,axis=0)
In [34]: Y
Out[34]: 
array([[0, 1, 2, 3],
       [0, 1, 2, 3]])
In [35]: Y[mask]
Out[35]: array([1, 3, 1, 2])

我們可以將其重塑為 2d - 如果每行的真數是一致的。

repeat的替代方法是：

In [39]: np.broadcast_to(y,mask.shape)[mask]
Out[39]: array([1, 3, 1, 2])

我希望這可以節省 memory，但速度並不快：

In [40]: timeit np.broadcast_to(y,mask.shape)[mask]
13.2 µs ± 335 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [41]: timeit y[None,:].repeat(2,axis=0)[mask]
4.72 µs ± 123 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

broadcast_to方法的擴展性可能更好，但我們只能通過測試來判斷。

有了where ，我們可以這樣做：

In [42]: np.nonzero(mask)
Out[42]: (array([0, 0, 1, 1]), array([1, 3, 1, 2]))
In [43]: y[np.nonzero(mask)[1]]
Out[43]: array([1, 3, 1, 2])