遮罩二维数组保留形状

Question

我有这样的二维 numpy 数组：

arr = np.array([[1,2,4],

                [2,1,1],

                [1,2,3]])

和一个布尔数组：

boolarr = np.array([[True, True, False],

                    [False, False, True],

                    [True, True,True]])

现在，当我尝试根据 boolarr 对 arr 进行切片时，它给了我

arr[boolarr]

输出：

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

但我希望有一个二维数组输出。所需的输出是

[[1, 2],

 [1],

 [1, 2, 3]]

Answer 1

使用的一个选项numpy是首先在 中添加行mask：take = boolarr.sum(axis=1)#array([2, 1, 3])然后像您一样屏蔽数组：x = arr[boolarr]#array([1, 2, 1, 1, 2, 3])并用于np.split根据np.cumsumof拆分平面数组take（因为函数需要拆分数组的索引）：np.split(x, np.cumsum(take)[:-1])[array([1, 2]), array([1]), array([1, 2, 3])]通用解决方案def mask_nd(x, m):    '''    Mask a 2D array and preserve the    dimension on the resulting array    ----------    x: np.array       2D array on which to apply a mask    m: np.array        2D boolean mask      Returns    -------    List of arrays. Each array contains the    elements from the rows in x once masked.    If no elements in a row are selected the     corresponding array will be empty    '''    take = m.sum(axis=1)    return np.split(x[m], np.cumsum(take)[:-1])例子让我们看一些例子：arr = np.array([[1,2,4],                [2,1,1],                [1,2,3]])boolarr = np.array([[True, True, False],                    [False, False, False],                    [True, True,True]])mask_nd(arr, boolarr)# [array([1, 2]), array([], dtype=int32), array([1, 2, 3])]或者对于以下数组：arr = np.array([[1,2],                [2,1]])boolarr = np.array([[True, True],                    [True, False]])mask_nd(arr, boolarr)# [array([1, 2]), array([2])]

Answer 2

您想要的输出不是二维数组，因为每个“行”都有不同数量的“列”。通过itertools.compress以下方式可以实现功能性非矢量化解决方案：from itertools import compressres = list(map(list, map(compress, arr, boolarr)))# [[1, 2], [1], [1, 2, 3]]

Answer 3

您可能正在寻找像掩码数组这样简单的东西。您可以使用掩码创建一个数组来屏蔽所需的值，这样它们就不会受到进一步操作的影响，也不会影响计算结果：marr = np.ma.array(arr, mask=~boolarr)请注意，必须翻转掩码，因为它是被掩码的无效元素。结果看起来像masked_array(data=[        [ 1  2 --]        [-- --  1]        [ 1  2  3]],    mask=[        [False False  True]        [ True  True False]        [False False False]],    fill_value = 999999)

Answer 4

In [183]: np.array([x[y] for x,y in zip(arr, boolarr)])Out[183]: array([array([1, 2]), array([1]), array([1, 2, 3])], dtype=object)应该在速度上有竞争力。（如果我们省略np.array外包装，只返回一个数组列表，速度会快一点。）但是需要进行实际的时间测试才能确定。