1305. 两棵二叉搜索树中的所有元素

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def getAllElements(self, root1: TreeNode, root2: TreeNode) -> List[int]:
        def visit(node: TreeNode):
            nodes, cur = [], node

            while cur or nodes:
                while cur:
                    nodes.append(cur)
                    cur = cur.left

                node = nodes.pop()
                yield node.val

                cur = node.right

        seq1, seq2 = visit(root1), visit(root2)

        ret = []
        i, j = None, None
        try:
            while True:
                i = next(seq1) if i is None else i
                j = next(seq2) if j is None else j
                if i < j:
                    ret.append(i)
                    i = None
                elif j < i:
                    ret.append(j)
                    j = None
                else:
                    ret.append(i)
                    ret.append(j)
                    i, j = None, None

        except StopIteration:
            pass
        finally:
            if i is not None:
                ret.append(i)
            if j is not None:
                ret.append(j)

        try:
            while True:
                ret.append(next(seq1))
        except StopIteration:
            pass

        try:
            while True:
                ret.append(next(seq2))
        except StopIteration:
            pass

        return ret