如何优雅地实现带掩码的迭代器配对（zip with mask）

碧海醫心 2025-12-29 00:00:00 次阅读

本文介绍一种更简洁、高效且符合 python 惯用法的方式，使用 itertools.compress 和 itertools.repeat 实现带布尔掩码或函数掩码的双序列配对，避免手动迭代与状态管理，提升可读性与健壮性。

在数据处理中，我们常需将一个较短的序列（如特征值 [1, 2, 3]）与一个较长的基准序列（如时间戳 [4, 5, 6, 7, 8, 9]）按条件对齐——例如仅在基准元素满足某条件（x >= 7）时才“激活”对应位置的特征值，其余位置填充 None。原始实现依赖显式迭代器控制和多重断言，逻辑复杂、易出错且难以复用。

更优解是借助标准库中的高阶工具组合实现声明式逻辑：

from itertools import compress, repeat, chain

def zip_mask(a, b, mask):
    """
    将可迭代对象 a 与 b 按 mask 对齐：
      - 若 mask 是 callable，作用于 b 的每个元素，返回布尔值；
      - 若 mask 是可迭代布尔序列，长度应与 b 一致。
    返回 zip 结果：(a_value or None, b_value)，其中 a_value 仅出现在 mask 为 True 的位置。
    """
    b_iter = iter(b)

    # 提取 b 中被 mask 选中的子序列
    if callable(mask):
        selected_b = compress(b_iter, (mask(x) for x in b))
    else:
        selected_b = compress(b_iter, mask)

    selected_b_list = list(selected_b)
    n_selected = len(selected_b_list)
    n_total = len(b)

    # 构造对齐后的 a 序列：前 (n_total - n_selected) 个为 None，后 n_selected 个为 a 的元素
    # 注意：要求 len(a) == n_selected，否则逻辑不成立（a 应恰好匹配被选中的 b 位置）
    if len(a) != n_selected:
        raise ValueError(f"Length mismatch: expected {n_selected} elements in 'a', got {len(a)}")

    padded_a = chain(repeat(None, n_total - n_selected), a)

    return zip(padded_a, b)

✅ 关键改进点：

使用 compress() 替代手动循环 + if mask(...) 判断，语义清晰、性能更优；
用 repeat(None, n) + chain() 构造稀疏对齐序列，避免索引错误与状态维护；
显式校验 len(a) == sum(mask)，提前暴露逻辑矛盾，而非静默错误；
函数签名与行为更贴近 itertools 风格：惰性生成、一次遍历、无副作用。

? 使用示例：

# 基于函数掩码
result = list(zip_mask([1, 2, 3], [4, 5, 6, 7, 8, 9], lambda x: x >= 7))
print(result)
# 输出：[(None, 4), (None, 5), (None, 6), (1, 7), (2, 8), (3, 9)]

# 基于布尔掩码列表
result2 = list(zip_mask(['x', 'y'], ['a', 'b', 'c', 'd'], [False, False, True, True]))
print(result2)
# 输出：[(None, 'a'), (None, 'b'), ('x', 'c'), ('y', 'd')]

⚠️ 注意事项：