All Elements in Two Binary Search Trees

Intuition

Solution

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def getAllElements(self, root1: TreeNode, root2: TreeNode) -> List[int]:
        inorder1 = []
        inorder2 = []
        self._dfs_inorder(root1, inorder1)
        self._dfs_inorder(root2, inorder2)

        # Merge 2 in-orders
        ans = []
        i1, i2 = 0, 0
        while i1 < len(inorder1) or i2 < len(inorder2):
            if i1 >= len(inorder1):
                ans.append(inorder2[i2])
                i2 += 1
            elif i2 >= len(inorder2):
                ans.append(inorder1[i1])
                i1 += 1
            elif inorder1[i1] > inorder2[i2]:
                ans.append(inorder2[i2])
                i2 += 1
            else:
                ans.append(inorder1[i1])
                i1 += 1
        return ans

    def _dfs_inorder(self, node, inorder):
        if node is None:
            return
        self._dfs_inorder(node.left, inorder)
        inorder.append(node.val)
        self._dfs_inorder(node.right, inorder)