python - 用掩码合并两个列表

Question

如果有一个布尔列表两个列表

两个列表长度之和等于掩码长度有没有一种有效的方法来合并这两个列表？

例子

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

直觉方式

 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

但它的速度很慢。我怎样才能提高速度？

标签： pythonnumpy

Answer 1

您可以使用数组和掩码来做到这一点：

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