Разделение маски на другие маски по ее элементам

Question

У меня есть маска массива python с произвольным размером, которая включает (0,1,2,3) числа.Я хочу разделить его на 4 подмаски, чтобы каждое из них содержало логические значения, первое из которых имеет значение True в позициях с 0 элементами, второе - True в положении с 1 элементом и так далее.Мой входной массив слишком большой, и мне нужен быстрый способ сделать это в Python.Вот пример:

mask = [[0 1 1]
        [2 0 3]
        [3,3,2]
        [1,1,1]]

Я хочу иметь:

sub_mask[0] = [[1 0 0]
               [0 1 0]
               [0 0 0]
               [0 0 0]]
sub_mask[1] = [[0 1 1]
               [0 0 0]
               [0 0 0]
               [1 1 1]]
sub_mask[2] = [[0 0 0]
               [1 0 0]
               [0 0 1]
               [0 0 0]]
sub_mask[3] = [[0 0 0]
               [0 0 1]
               [1 1 0]
               [0 0 0]]

Answer 1

Вы можете сделать это с помощью простого понимания списка:

sub_mask = [
    mask == val
    for val in np.unique(mask)
]

Answer 2

Вы можете использовать numpy:

import numpy as np

mask = [[0, 1, 1],
        [2, 0, 3],
        [3, 3, 2],
        [1, 1, 1]]
mask = np.array(mask)

def get_sub_mask(n):
    # Create an empty sub-mask with all 0's
    sub_mask = np.zeros(mask.shape)
    # Set elements to 1 based on their position in the original mask
    sub_mask[np.where(mask == n)] = 1
    return sub_mask