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

Question

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

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

我试图使用掩码来获取每一行的元素。所以输出应该是这样的：

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

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

有什么想法可以将其作为2D数组返回吗？

Answer 1

怎么样

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