沿着任意軸的具有邏輯值的屏蔽陣列

Question

假設我有一個多維數組和一個邏輯值向量。我想選擇沿任意（第n）維的項目。在以下示例中，我要沿第二維度選擇第一和第三值：

>>> A = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) 
>>> mask = np.array([True, False, True, False]) 
>>> dim_to_mask = 1 # i.e. 2nd dimension because it's 0-indexed 
>>> B = ... # here do mask the dim_to_mask-th dimension - HOW??? 
>>> B 
[[1, 3], 
[5, 7], 
[9, 11]] 

注意：假設邏輯向量的長度對應於給定軸線的長度。

我知道如果數組只是使用[]運算符的一維數據，那很容易，但這是多維問題。

其實我想要的東西像功能take(indices, axis)它沿任意軸選擇給定的指標。唯一的區別是我確實有邏輯值而不是數字索引。

我也瞄準最快的解決方案，以便將向量的邏輯值轉換爲索引，並使用take可能不是最好的解決方案。

我想這一定是明顯的，我失蹤了。 :)

Answer 1

你可以使用np.compress：

>>> A.compress(mask, axis=1) 
array([[ 1, 3], 
     [ 5, 7], 
     [ 9, 11]]) 

此功能沿特定軸返回數組的切片。它接受一個布爾數組來進行選擇。