logical masking in Python

Question

In Matlab, it is natural to make heavy use of the built-in logical indexing and masking mechanisms. Eg, (1) idx = (A == 2) creates a logical array the same size as A, identifying all elements equal to 2, and (2) C = B(idx) pulls all the corresponding elements out of B, assuming sizes of A and B are the same.

This mechanism is elegant and extremely efficient. It works regardless of the dimensionality of A & B. In Matlab, it is more efficient to use the logical mask than to extract the indices numerically. I've just started learning Python, which does everything differently. What would be the closest equivalent to (1) and (2) in standard Python 3 (without installing NumPy, etc)?

Answer 1

This applies a function to all elements of a collection. The built-in ways to do this are map and comprehensions (map(lambda x: x==2, A), [x==2 for x in A]). You would have to do your own extension of map if you want to apply a function to randomly nested lists.

def is_two(x):  # more readable than the lambda
    return x == 2

def map_nested(fnc, lst):
    if not isinstance(lst, list):
        return fnc(lst)
    return [map_nested(fnc, x) for x in lst]

A = [[1,2,3], [2,3,4], 1, 2]
map_nested(is_two, A)
#[[False, True, False], [True, False, False], False, True]