Remove all elements in each list of a nested list, based on first nested list

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]]

>Solution :

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]]

Try it online!

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 is falsy (0 here), the corresponding element is dropped, when it’s truthy (1 here) it’s retained.