Select Columns by Masking Using Itertools Compress and a Filter List


In this example we use compress from itertools to select columns in a dataframe by using a filter list of boolean values as a mask. Here we've manually defined the filter list but this could be derived in other ways such as through a recursive feature elimination algorithm.

The original dataframe has four columns but only the 1st, 3rd and 4th columns will be kept.

 1|  from itertools import compress
 3|  filter_list = [True, False, True, True]
 4|  cols = compress(df.columns, filter_list)
 5|  cols = list(cols)
 6|  df.columns = cols
Did you find this snippet useful?

Sign up for free to to add this to your code library