Selecting subset of elements of an array/list with boolean indices-CodePudding

I have a list (or numpy array), let's call it my_list with length of 100. I have another list, called flag_list with the same length consisting of 1s and 0s. I want my new_list be the elements of my_list whom have 1 in their corresponding position in flag_list. Somehow, to remove those who have 0 flag.

Is there a short way of doing this without writing a for loop?

I also have a 2d numpy array that I want to exclude some raws based on flags.

Thanks in advane

CodePudding user response：

You can use list comprehension and enumerate to accomplish this

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)

Outputs

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

CodePudding user response：

Boolean indexing can be used for this. Convert list of 0,1s to array of booleans and then do indexing.

import numpy as np

my_list = np.arange(10)
flag_list = np.random.randint(2,size=(1,10))
flag_list = flag_list[0]
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])

Output:

  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]