Маска 2D массива с сохранением формы

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

Теперь, когда я пытаюсь нарезать arr на основе boolarr, он дает мне

arr[boolarr]

Вывод:

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

Но вместо этого я хочу получить вывод двумерного массива.Желаемый вывод

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

Answer 1

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 обертку, возвращая только список массивов.)

Но для уверенности необходимы реалистичные тесты времени.

Answer 2

Возможно, вы ищете что-то столь же простое, как массив маски .Вы можете использовать маску для создания массива, который маскирует нужные значения, чтобы они не влияли на дальнейшие операции и не влияли на результаты вычислений:

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 3

Ваш желаемый результат не является двумерным массивом, поскольку каждая «строка» имеет различное количество «столбцов».Функциональное не векторизованное решение возможно через itertools.compress:

from itertools import compress

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

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

Answer 4

Опция, использующая 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 5

Вот один из способов сделать это с помощью list:

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