掩模2D阵列保持形状

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

但我希望有一个二维数组输出，而不是。所期望的输出是

[[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.cumsum的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

所需输出不是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包裹，只返回一个数组列表。）

但现实的时间测试需要确定。