选择具有布尔索引的数组/列表的元素子集

Question

我有一个列表（或 numpy 数组），我们称它为 my_list，长度为 100。我有另一个列表，称为flag_list，长度相同，由 1 和 0 组成。我希望我的new_list 是my_list 的元素，它们在flag_list 中的相应位置有1。不知何故，删除那些有 0 标志的人。

有没有不写 for 循环的捷径？

我还有一个 2d numpy 数组，我想根据标志排除一些原始数据。

提前感谢

Answer 1

您可以使用list comprehension 和enumerate 来完成此操作

my_list = list('abcde')
flag_list = [0,1,0,1,1]
new_list = [val for i, val in enumerate(my_list) if flag_list[i]]
print(new_list)

输出

['b', 'd', 'e']

Answer 2

布尔索引可用于此。将 0,1s 列表转换为布尔数组，然后进行索引。下面的解决方案通过简单的 numpy 操作完成了这项工作，没有任何高阶函数或理解。

import numpy as np

# generate dummy data
my_list = np.arange(10)
flag_list = np.random.randint(2,size=(10,))

print("  my_list",  my_list)
print("flag_list", flag_list)

bool_list = flag_list == 1
print("bool_list", bool_list)
print("my_list[bool_list]", my_list[bool_list])

输出：

  my_list [0 1 2 3 4 5 6 7 8 9]
flag_list [0 0 1 1 1 0 0 1 1 0]
bool_list [False False  True  True  True False False  True  True False]
my_list[bool_list] [2 3 4 7 8]

如果flag_list是一个普通的python列表而不是一个由0和1组成的numpy数组，那么我们可以直接转换为bool数组

bool_list = np.array(flag_list, dtype='bool')

Answer 3

尝试使用压缩方式解决问题：

my_list = ['a','b','c','d','e']
flag_list = [0,0,1,0,1]
res = [i for i,j in zip(my_list,flag_list) if j==1]
print(res)

输出：

['c', 'e']

使用过滤方法：

my_list = ['a','b','c','d','e']
flag_list = [0,0,1,0,1]
lst = list(filter(lambda x:x[1]==1,zip(my_list,flag_list)))
res = [i[0] for i in lst]
print(res)

输出：

['c', 'e']