如何将Haskell惰性countRuns函数移植到Python并解决迭代器冲突问题

核心问题原因 #

• 代码漏了将输入可迭代对象转为迭代器、取出首个元素x的步骤，原代码中的x属于未定义变量，运行会直接报错。
• Python迭代器是单向消耗的，你将同一个迭代器同时传给takewhile和dropwhile时，takewhile会遍历到第一个不符合条件的元素才停止，这个不符合条件的元素会被takewhile消耗，后续dropwhile会从该元素之后开始遍历，直接丢失每组开头的非重复元素。

修复后的代码 #

from itertools import takewhile, dropwhile, tee

def count_runs(xs):
    # 统一将输入转为迭代器，兼容所有可迭代类型
    it = iter(xs)
    try:
        # 取出首个元素，是获取可迭代对象首元素的通用方式
        x = next(it)
    except StopIteration:
        # 迭代器为空，终止递归
        return
    # 拆分出两个独立的迭代器副本，避免takewhile和dropwhile互相消耗
    it_for_take, it_for_drop = tee(it)
    us = takewhile(lambda y: y == x, it_for_take)
    vs = dropwhile(lambda y: y == x, it_for_drop)
    yield (1 + len(list(us)), x)
    yield from count_runs(vs)

效果验证 #

运行你给出的测试用例：

list(count_runs(['a', 'a', 'b', 'c', 'a', 'd', 'd'])) 
# 输出：[(2, 'a'), (1, 'b'), (1, 'c'), (1, 'a'), (2, 'd')]

完全符合预期。

补充说明 #

如果需要处理连续组数量极多的场景，Python默认递归深度限制（通常为1000）可能导致栈溢出，可改用迭代版本实现，性能和稳定性更高：

def count_runs(xs):
    it = iter(xs)
    try:
        current = next(it)
    except StopIteration:
        return
    count = 1
    for item in it:
        if item == current:
            count += 1
        else:
            yield (count, current)
            current = item
            count = 1
    yield (count, current)

内容的提问来源于stack exchange，提问作者Eric Auld