python从布尔列掩码创建数组的子集

Question

我需要创建一个能够改变这个输入的函数或方程......

a = [True, False, True]
b = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

进入此输出...

c = [[1, 3], [4, 6], [7, 9]]

注意：上面的数组都是numpy数组。谢谢！

Answer 1

numpy解决方案

import numpy as np 
a = np.array([True, False, True])
b = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
b[:,a]

array([[1, 3],
       [4, 6],
       [7, 9]])

Answer 2

您可以使用itertools compress

from itertools import compress
c = [list(compress(i, a)) for i in b]

[[1, 3], [4, 6], [7, 9]]