如何使用二维数组掩码索引 numpy 二维数组

Question

假设我们有一个 numpy 二维数组，如下所示：

x = array([[0, 7, 1, 6, 2, 3, 4],
           [4, 5, 0, 1, 2, 7, 3]])

和如下所示的 2D 蒙版：

mask = array([[ True, False,  True, False, False, False, False],
              [False, False, False, False,  True, False, False]])

我正在尝试使用掩码来获取每一行的元素。 所以 output 应该看起来像这样：

array(
    [0, 1],
    [2]
)

如果我使用x[mask]我得到array([0, 1, 2])这是错误的，因为它展平了所有选定的项目。

有什么想法可以将它作为二维数组返回吗？

Answer 1

怎么样

[xi[mi] for xi,mi in zip(x,mask)]