I have a following list:
a = [[0, 1, 0, 1, 1, 1], [23,22, 12, 45, 32, 33],[232, 332, 222, 342, 321, 232]]
I want to remove 0
in a[0]
and corresponding values of a[1]
and [2]
, so the result list should be as follows:
d = [[1, 1, 1, 1], [22, 45, 32, 33], [332, 342, 321, 232]]
Advertisement
Answer
itertools.compress
is built for this task. Pair it with a listcomp to operate on each sublist, and force it to eagerly produce each new sublist, and you get what you want with fairly simple code:
from itertools import compress a = [[0, 1, 0, 1, 1, 1], [23,22, 12, 45, 32, 33],[232, 332, 222, 342, 321, 232]] new_a = [list(compress(sublst, a[0])) for sublst in a] print(new_a)
Which outputs:
[[1, 1, 1, 1], [22, 45, 32, 33], [332, 342, 321, 232]]
Each call uses a[0]
as the selectors to determine which elements to keep from the data (each new sub-list
); when the selector value from a[0]
is falsy (0
here), the corresponding element is dropped, when it’s truthy (1
here) it’s retained.