Question

我有2D 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])

但是我正在寻找2D数组输出。所需的输出是

[[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根据take的{{3}}拆分平面数组（因为该函数期望索引在哪里拆分数组）：

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

您想要的输出不是2D数组，因为每个“行”都有不同数量的“列”。通过itertools.compress可以实现功能非矢量化的解决方案：

from itertools import compress

res = list(map(list, map(compress, arr, boolarr)))

# [[1, 2], [1], [1, 2, 3]]

Answer 3

以下是使用list进行此操作的一种方法：

[[arr[row][col] for col in range(3) if boolarr[row][col]] for row in range(3)]
# [[1,2], [1], [1,2,3]]

Answer 4

您可能正在寻找像masked array这样简单的内容。您可以使用遮罩创建一个遮罩所需值的数组，以使它们不受进一步操作的影响，并且不影响计算结果：

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 5

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包装，只返回数组列表，则速度会更快一些。）

但是需要进行实际的时间测试。

遮罩2D阵列保留形状

5 个答案: