Reputation: 367
I have a python array mask with arbitrary size which includes (0,1,2,3) numbers. I want to divide it into 4 sub-masks that every one of them includes boolean values which the first one has True in the 0-element positions, second one has True in 1-element position and so on. My input array is too large and I need a fast way to do this in python. Here is an example:
mask = [[0 1 1]
[2 0 3]
[3,3,2]
[1,1,1]]
I want to have:
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]]
Upvotes: 0
Views: 326
Reputation: 5470
You can do this via a simple list comprehension:
sub_mask = [
mask == val
for val in np.unique(mask)
]
Upvotes: 2
Reputation: 1614
You can use 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
Upvotes: 2