沿任意轴具有逻辑值的掩码数组

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

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