How to mask 2D list in Python?

Question

I created a 2D list and applied a mask on the 2D list followed the same type question. It turned out the posted solution simply didn't work for 2D list. Here is the code and output:

from itertools import compress
class MaskableList(list):
    def __getitem__(self, index):
        try: return super(MaskableList, self).__getitem__(index)
        except TypeError: return MaskableList(compress(self, index))

aa=[['t', 'v'], ['a', 'b']]
aa=MaskableList(aa)
print(aa)
>>> [['t', 'v'], ['a', 'b']]

mask1=[[1,0],[0,1]]
print(aa[mask1])
>>> [['t', 'v'], ['a', 'b']]

mask2=[1,0,0,1]
print(aa[mask2])
>>> [['t', 'v']]

It there a clean and efficient way which works for masking 2D list.

Answer 1

A simple way you could do it would be to reimplement the generator expression that itertools.compress implements.

You turn it into statements, and when at a given position both the data and the selector are lists, you recurse that new compress function over that sublist:

from collections import Iterable
def compress_rec(data, selectors):
    for d, s in zip(data, selectors): # py3 zip is py2 izip, use itertools.zip_longest if both arrays do not have the same length
        if isinstance(d, Iterable) and isinstance(s, Iterable):
            yield compress_rec(d, s)
        else:
            yield d

that way it would work with any dimensional array.

HTH

Answer 2

There is definitely a better solution that involves messing with the class definition but a workaround would be this:

from itertools import compress
class MaskableList(list):
    def __getitem__(self, index):
        try:
            return super(MaskableList, self).__getitem__(index)
        except TypeError:
            return MaskableList(compress(self, index))

aa = [['t', 'v'], ['a', 'b']]
mask1 = [[True, False], [False, True]]

new = [MaskableList(sublist)[submask] for sublist, submask in zip(aa, mask1)]
print(new)  # -> [['t'], ['b']]