根据另一个 0/​​1 索引数组从 numpy 数组中提取值

Question

给定一个仅包含 0 和 1 个元素的索引数组idx，1 代表感兴趣的样本索引，以及一个样本数组A( A.shape[0] = idx.shape[0])。这里的目标是根据索引向量提取样本子集。

在 matlab 中，执行以下操作很简单：

B = A(idx,:) %assuming A is 2D matrix and idx is a logical vector

如何在Python中以简单的方式实现这一点？

Answer 1

如果您的掩码数组idx与数组具有相同的形状，那么如果使用 转换为布尔数组，A您应该能够提取掩码指定的元素。idxastype

演示 -

>>> A
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])
>>> idx
array([[1, 0, 0, 1, 1],
       [0, 0, 0, 1, 0],
       [1, 0, 0, 1, 1],
       [1, 0, 0, 1, 1],
       [0, 1, 1, 1, 1]])

>>> A[idx.astype(bool)]
array([ 0,  3,  4,  8, 10, 13, 14, 15, 18, 19, 21, 22, 23, 24])