CODEWITHSUNDEEP
pythonnumpy
yhChou
yhChou

Reputation: 23

Merging two list with a mask

If have a boolean list two lists

The sum of two lists length is equal to the mask length Is there an efficient way to merge this two list?

Example

 mask   : array([False, True, True, False, False])
 list a : array([3, 4])
 list b : array([1, 2, 5])

 if True fill in a if False fill in b 
 output : array([1, 3,  4, 2, 5])

Intuition way

 a_count = 0
 b_count = 0
 output = np.zeros(mask.shape)
 for i in range(len(mask)):
   if mask[i] == True:
     output[i] = a[a_count]
     a_count += 1
   elif mask[i] == False:
     output[i] = b[b_count]
     b_count += 1

But it's to slow. How can I improve the speed?

Upvotes: 1

Views: 168

Answers (1)

Paddy Harrison
Paddy Harrison

Reputation: 2002

You can do it with arrays and masking:

mask = np.array([False, True, True, False, False])
a = np.array([3, 4])
b = np.array([1, 2, 5])

# create output array
out = np.empty_like(mask, dtype=int)

# select elements from out that are True in mask -> set them to a
out[mask] = a
# if False set to b 
out[~mask] = b

out
>>> array([1, 3, 4, 2, 5])

Upvotes: 2

Related Questions