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

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

此函数返回数组沿特定轴的切片。它接受一个 bool 数组来进行选择。