有没有更好的方法用掩码压缩可迭代对象？

Question

我需要使用个人掩码功能或仅掩码列表来压缩可迭代对象

[True, False, ...]

from collections.abc import Callable

def zip_mask(a, b, mask):
    iter_a, iter_b = iter(a), iter(b)
    if isinstance(mask, Callable):
        assert len(a) == (valid := sum(mask(i) for i in b)), f'have mismatch length with `{len(a)}`; `{valid}`'
        for i in b:
            yield next(iter_a) if mask(i) else None, next(iter_b)
    else:
        assert (select := set(mask)) == {1, False}, f'bad input mask selections {select}'
        assert len(b) == len(mask) > len(a), f'bad input shapes'
        assert len(a) == (valid := sum(mask)), f'have mismatch length with `{len(a)}`; `{valid}`'
        for m in mask:
            yield next(iter_a) if m else None, next(iter_b)


# (None, 4) (None, 5) (None, 6) (1, 7) (2, 8) (3, 9)
print(*zip_mask([1, 2, 3], [4, 5, 6, 7, 8, 9], lambda x: x >= 7))

希望有更好的方法来做到这一点。

Answer 1

from itertools import compress, repeat

def zip_mask(a, b, mask):
    iter_b = iter(b)
    if callable(mask):
        filtered_b = compress(iter_b, (mask(i) for i in iter_b))
    else:
        filtered_b = compress(iter_b, mask)

    filtered_b_list = list(filtered_b)  # Convert to list for length calculation

    # Ensure both parts are iterables for concatenation
    yield from zip(list(repeat(None, len(b) - len(filtered_b_list))) + a, b)
    
result = list(zip_mask([1, 2, 3], [4, 5, 6, 7, 8, 9], [False, False, False, True, True, True]))
print(result)