How to use `filter` with multiple iterables, as is supported by `map`?

Viewed 545

filter accepts only one iterable, whereas map accepts a variadic number of iterables. For example, I can exhaust map(operator.add, [1, 2, 3, 4], [1, 2, 2, 4]) to get [2, 4, 5, 8].

I'm looking for a similar mechanism for filter, accepting any predicate and a variable number of iterables. Exhausting filter(operator.eq, [1, 2, 3, 4], [1, 2, 2, 4]) causes a TypeError about how filter only accepts 1 iterable, not 2.

My expected output for that particular case is ([1, 2, 4], [1, 2, 4]), i.e the pairwise elements that don't satisfy operator.eq are removed.

Here's what I have so far (eager version supporting only 2 iterables instead of N):

from typing import TypeVar, Callable, Iterable

A = TypeVar("A")
B = TypeVar("B")

def filter_(predicate: Callable[[A, B], bool], iterable1: Iterable[A], iterable2: Iterable[B]) -> (Iterable[A], Iterable[B]):
    filtered_iterable1 = []
    filtered_iterable2 = []

    for value1, value2 in zip(iterable1, iterable2):
        if predicate(value1, value2):
            filtered_iterable1.append(value1)
            filtered_iterable2.append(value2)

    return filtered_iterable1, filtered_iterable2

However my goal is to 1) be able to support N iterables and 2) to have filter_ be lazy instead of eager, as is with filter.

3 Answers

Unfortunately there's no equivalent to starmap like starfilter, so the equivalent I can think of is:

[i for i in zip(*lists) if predicate(*i)]

lists here being something like ([..], [..]). This results in:

[(1, 1), (2, 2), (4, 4)]

To turn this back into separate lists, use tuple(map(list, zip(*result))):

([1, 2, 4], [1, 2, 4])

So, putting it together:

predicate = operator.eq
lists = [1, 2, 3, 4], [1, 2, 2, 4]

result = tuple(map(list, zip(*(i for i in zip(*lists) if predicate(*i)))))

Your answer is in your implementation. Map accepts a function taking multiple lists which must match the number of arguments. Filter takes a single list to filter, so the difference in not just semantic - it makes sense for filter to only take a single list. In your case the list is indeed the zip, and that is what you implement. What you are missing is a nifty way to unpair the paired results:

>>> r1, r2 = zip(*filter(lambda x: predicate(*x), zip([1, 2, 3, 4, 5], [1, 1, 3, 3, 5)))
>>> r1
(1, 3, 5)
>>> r2
(1, 3, 5)

How about:

def filter_(predicate, *iterables):
    for t in zip(*iterables):
        if predicate(*t):
            yield t

print(list(filter_(operator.eq, [1, 2, 3, 4], [1, 2, 2, 4])))

It is lazy, it outputs [(1, 1), (2, 2), (4, 4)] for your test case, and no, you can't have ([1, 2, 4], [1, 2, 4]) as a result in a lazy way. To convert from [(1, 1), (2, 2), (4, 4)] to ([1, 2, 4], [1, 2, 4]) you could use: zip(*filter_(operator.eq, [1, 2, 3, 4], [1, 2, 2, 4])) but then of course you lose the laziness.

Related