沿任意轴屏蔽具有逻辑值的数组

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]])

此函数返回沿特定轴的数组切片。它接受一个布尔数组，用于进行选择。