遮罩二维数组保留形状

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 compress

res = 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外包装，只返回一个数组列表，速度会快一点。）

但是需要进行实际的时间测试才能确定。