python迭代器实现LeetCode No.100 相同的树

简介

• 今天《fluent python》正好看完迭代器/生成器/协程部分
• 发现里边的一些知识很适合用到这道题
• getgeneratorstate()用于获取生成器的状态，如果迭代完成后，两生成器都处于GEN_CLOSED状态则说明两树相同
• 这种生成器实现是惰性的，也能实现两树的边遍历边比较

from inspect import getgeneratorstate


class Solution:
    def isSameTree(self, p: TreeNode, q: TreeNode) -> bool:
        if (not p) and (not q):
            return True
        elif not p:
            return False
        elif not q:
            return False
        
        gen_p, gen_q = self.tree_iter(p), self.tree_iter(q)
        
        for i, j in zip(gen_p, gen_q):
            if i != j:
                return False

        # zip()在其中一个迭代器返回StopIteration时就会返回StopIteration来结束for循环
        # 手动调用next()来使两迭代器迭代次数匹配
        try:
            next(gen_q)
        except StopIteration:
            return getgeneratorstate(gen_p) == getgeneratorstate(gen_q)
        
    # 实现前序遍历迭代器
    def tree_iter(self, tree):
        stack = []
        p = tree
        while p or stack :
            while p:
                yield p.val
                stack.append(p)
                p = p.left
            else:
                yield None
            node = stack.pop()
            p = node.right

总结（关于zip函数的使用细节）

• 通过zip()得到的生成器在其中一个子生成器抛出StopIteration异常时就会向上抛出StopIteration异常，从而结束for循环
• 此时另外的生成器并未达到GEN_CLOSED状态